Agent模块化开发之Modular RAG应用
·
内容提要
- 模块化开发概述
- Modular RAG应用
一、模块化开发概述
我们在进行Agent应用开发时，需要按业务(功能)分工不同，划分不同模块，分别进行开发。
比如：搜索模块、数据解析、问答模块、配置模块等
1.1、模块化开发重要性
模块化是实现职责分离、提升可维护性与可扩展性，让AI Agent系统更高效、稳定地运行。
在Agent应用开发中，模块化设计将复杂任务拆解为多个独立功能单元，例如规划、执行、记忆、工具调用等。每个模块专注处理特定职责：
- 规划模块，负责任务分解与路径设计；
- 执行模块，生成并运行代码；
- 记忆模块，存储和检索历史信息（存储关注度内容）；
- 工具模块，集成外部API或服务能力；
- 上下文管理，确保信息流动连贯（实现承上启下的内容）；
- 反思模块，进行错误检测与自我修正。
1.2、模块化结构带来三大优势
- 降低耦合度：模块间通过清晰接口通信，单个模块变更不会波及整体系统；
- 增强可复用性：通用功能（如API调用、日志记录）可封装成独立模块，在不同项目中重复使用；
- 便于调试与优化：问题定位更精准，性能瓶颈可针对性改进。
二、Modular RAG概述
2.1、Modular RAG简介
1、什么是Modular RAG？
Modular RAG(模块化的RAG)是一种将 RAG 系统拆分为独立、可插拔的模块化架构设计的技术。与传统的 RAG（向量检索 + LLM 生成）不同，Modular RAG 允许灵活组合各种检索器、处理器、重排序器、融合模块(不组成部分)等。
它在Agent应用的模块开发中应用非常广泛。
传统 RAG: 查询 → 向量检索 → 大模型 → 答案
Modular RAG: 查询 → [查询改写] → [多路检索](bm25，向量检索) → [重排序] → [融合]->LLM → [生成] → 答案
2、核心模块
|
模块 |
功能 |
示例组件 |
|
Query 模块 |
查询理解与改写 |
查询扩展、HyDE(假定性文档)、多查询生成 |
|
Retriever 模块 |
多源检索 |
向量检索、关键词检索、数据库查询 |
|
Post-Processing 模块 |
检索后处理 |
重排序、去重、过滤、压缩（将拿到的相关文档） |
|
Fusion 模块 |
多路结果融合 |
RRF、加权融合、LLM融合，找较优内容的前TOP-N |
|
Generator 模块 |
答案生成 |
上下文构建、LLM调用 |
代表AI应用开发程序的核心组成部分，也指示开发的流程与步骤。
3、模块化设计的优势
- 可组合性 —— 根据场景自由选择和组合模块
- 可扩展性 —— 轻松添加新的检索器或处理器
- 可测试性 —— 每个模块独立测试、调试
- 可优化性 —— 针对特定模块进行专项优化
2.2、完整的多模块RAG系统代码实现
本RAG系统具备如下功能：演示多模块RAG应用开发的流程，实际开发替换模拟数据为真实数据即可。
✅ 查询改写（多角度生成）-》Query模块
✅ 中文分词和检索-》Retriever模块
✅ 多路检索（向量 + BM25）-》Retriever模块
✅ 重排序 -》Post-Processing 模块
✅ 上下文压缩 -》Post-Processing 模块
✅ RRF结果融合-》Fusion模块，找TOP-N
✅ 答案生成-》Generator 模块
环境准备：
# 需要 Python 3.8+
# 安装依赖:scikit-learn：python机器学习库，数据分析、预处理、模型开发和评估等
# transformers、torch深度学习模型（框架），用于复杂文本，图形图像等处理
###################### 需要安装的库 ##########################
pip install numpy scikit-learn transformers torch
###################### 需要在程序文件中导入 ###################
import numpy as np
from typing import List, Dict, Any, Optional, Callable, Tuple
from dataclasses import dataclass, field
from collections import defaultdict
import re
# python自带的哈希计算工具
import hashlib
项目结构参考：
2.2.1、基础数据结构
文件：my_query.py
# =============================================================================
# 第一部分：基础数据结构定义
# =============================================================================
# @dataclass 是 Python 3.7+ 引入的一个装饰器，来自 dataclasses 模块。它的核心作用是自动帮你生成类的样板代码（比如 __init__、__repr__、__eq__ 等方法），让你写数据容器类时更简洁。
# @dataclass装饰一个类时，Python 会根据你定义的类属性（类型注解）自动生成以下方法：
# __init__：构造函数
# __repr__：打印对象时的字符串表示
# __eq__：判断两个对象是否相等（基于字段值）
# __hash__（可选）：用于放入 set/dict 的 key（需显式开启）
from dataclasses import field, dataclass
from typing import Dict, Any
# 类数据结构的装饰器，用于自动生成 __init__、__repr__、__eq__ 等方法
# 节省代码，提高效率
# ... existing code ...
@dataclass
class Document:
"""
文档对象
"""
id: str # 文档ID
title: str # 标题
content: str # 内容
query: str = "" # 查询内容
metadata: Dict[str, Any] = field(default_factory=dict) # 内容的元数据
score: float = 0.0 # 评分
def __repr__(self):
return f"Document(id={self.id}, title={self.title}, content={self.content})"
def __eq__(self, other):
if not isinstance(other, Document):
return False
return self.id == other.id and self.title == other.title and self.content == other.content
def __hash__(self):
return hash((self.id, self.title, self.content))
@dataclass
class RetrievalResult:
"""
检索结果对象
"""
query:str # 待查询的查询内容
documents: list[Document] # 已检索到的文档列表
sources:str # 检索来源(msyql,vector)
# 测试一下
if __name__ == "__main__":
doc = Document(id="1", title="hello world", content="hello world", query="hello world")
print(doc)
print(doc.query)
print(doc.metadata)
print(doc.score)
doc2 = Document(id="2", title="兵王之战", content="因为一次意外任务，林凡返回都市...", query="霸道")
retrieval_result = RetrievalResult(query="争斗", documents=[doc, doc2], sources="mysql")
print(retrieval_result)
2.2.2、查询改写模块（Query Rewriting）
文件：query_rewriter.py
# new
# =============================================================================
# 第二部分：查询改写模块（Query Rewriting）
# =============================================================================
import re
from dataclasses import dataclass
from typing import List
@dataclass
class QueryRewriter:
"""
查询改写模块
为什么需要查询改写？
1. 用户输入的查询可能不够精准
2. 同义词扩展可以增加召回
3. 多角度查询可以覆盖不同方面
4. HyDE（假设性文档）可以从答案反推相关文档
支持的改写策略：
1. 同义词扩展：为查询词添加同义词
2. HyDE：生成一个假设性的理想答案
3. Multi-Query：生成多个不同视角的查询
"""
synonyms = {
"好": ["优秀", "出色", "棒", "卓越"],
"坏": ["差", "糟糕", "恶劣", "劣质"],
"大": ["巨大", "庞大", "宏大", "巨型"],
"小": ["微小", "小巧", "迷你", "紧凑"],
"快": ["快速", "迅速", "飞速", "迅猛"],
"慢": ["缓慢", "迟缓", "徐缓", "迟钝"],
"新": ["崭新", "新颖", "现代", "最新"],
"旧": ["古老", "陈旧", "传统", "老式"],
"重要": ["重大", "关键", "核心", "首要"],
"难": ["困难", "艰难", "复杂", "棘手"]
}
def expand_with_synonyms(self, query: str) -> List[str]:
"""
同义词扩展
原理：
对于查询中的每个词，如果存在同义词，生成几个包含同义词的新查询
示例：
输入: "好产品"
输出: ["好产品", "优秀产品", "出色产品"]
参数：
query: 原始查询
返回：
扩展后的查询列表
"""
words = query.lower().split()
expanded_queries = [query] # 包含原查询
# 对每个词尝试替换同义词
for i, word in enumerate(words):
if word in self.synonyms:
# 取前 2 个同义词，避免生成过多查询
for syn in self.synonyms[word][:2]:
new_words = words.copy()
new_words[i] = syn
expanded_queries.append(" ".join(new_words))
# 去重并限制数量
return list(set(expanded_queries))[:3]
def hyde_pseudo_doc(self, query: str) -> str:
"""
HyDE (Hypothetical Document Embeddings)【假设文档嵌入】
核心思想：
1. 让 LLM 根据查询生成一个"假设性的理想文档"
2. 这个假设文档应该包含答案所需的信息
3. 用这个假设文档去检索，可以找到更相关的真实文档
注意：这里使用模板模拟，实际应该调用 LLM
参数：
query: 原始查询
返回：
假设性文档文本
"""
templates = [
f"以下是关于{query}的详细信息：{query}是一个重要的话题，涉及多个方面。"
f"关于{query}的关键点包括其定义、应用场景和在该领域的重要意义。"
f"研究表明，{query}对理论和实践都有许多启示。",
f"本文档全面探讨了{query}。描述了{query}的主要特征。"
f"我们发现{query}与几个重要概念密切相关。"
f"理解{query}需要考虑多个角度。",
f"本分析聚焦于{query}。有证据表明{query}发挥着至关重要的作用。"
f"进一步的研究揭示{query}具有显著的影响力。"
f"这里总结了对{query}的研究发现。"
]
# 选择最合适的模板（基于查询长度？简单起见用第一个）
return templates[0]
def multi_query_generation(self, query: str) -> List[str]:
"""
多查询生成（优化版）
"""
queries = [query]
# 清理查询文本
query_clean = re.sub(r'[？?！!。，,、；;：:]', '', query)
# 策略1：变化句型
if "？" in query or "?" in query:
# 如果是疑问句，去掉问号
no_q = query.replace("？", "").replace("?", "").strip()
if len(no_q) >= 2: # 过滤太短的
queries.append(no_q)
else:
# 如果是陈述句，生成疑问版本
if len(query) >= 2:
queries.append(f"{query}是什么？")
# 策略2：提取关键词版本（取前3个有意义的词）
keywords = self.extract_keywords(query_clean)
if keywords and len(keywords) >= 2 and keywords != query:
queries.append(keywords)
# 过滤掉太短的查询（长度小于2的不要）
queries = [q for q in set(queries) if len(q.strip()) >= 2]
# 限制数量
return queries[:3]
def extract_keywords(self, text: str) -> str:
"""提取有意义的 keywords"""
# 去除英文和数字混合词中的标点
text = re.sub(r'[^\u4e00-\u9fff\w]', ' ', text)
words = text.split()
# 停用词
stopwords = set(['的', '了', '是', '在', '我', '有', '和', '就', '不', '人',
'都', '一', '一个', '上', '也', '很', '到', '说', '要', '去',
'你', '会', '着', '没有', '看', '好', '自己', '这', '那',
'什么', '如何', '为什么', '怎么', '哪个', '哪些'])
# 过滤停用词和短词
keywords = [w for w in words if w not in stopwords and len(w) >= 2]
# 如果提取不到关键词，返回原文本的前6个字符
if not keywords:
return text[:6] if len(text) > 2 else text
return " ".join(keywords[:3])
def tokenize_for_rewrite(self, text: str) -> List[str]:
"""为改写提取关键词"""
# 去除标点
text = re.sub(r'[^\u4e00-\u9fff]', ' ', text)
words = text.split()
# 过滤停用词
stopwords = set(
['的', '了', '是', '在', '我', '有', '和', '就', '不', '人', '都', '一', '一个', '上', '也', '很', '到',
'说', '要', '去', '你', '会', '着', '没有', '看', '好', '自己', '这', '那'])
return [w for w in words if w not in stopwords and len(w) >= 2]
def rewrite(self, query: str, method: str = "multi") -> List[str]:
"""
查询改写入口
参数：
query: 原始查询
method: 改写方法 ("expand", "hyde", "multi")
返回：
改写后的查询列表
"""
if method == "expand":
return self.expand_with_synonyms(query)
elif method == "hyde":
# HyDE：返回原查询 + 假设文档
return [query, self.hyde_pseudo_doc(query)]
elif method == "multi":
return self.multi_query_generation(query)
else:
return [query]
# 测试一下
if __name__ == "__main__":
qr = QueryRewriter()
print(qr.rewrite("What,.!is(the capital of France?", "expand"))
list=qr.multi_query_generation("What,.!is(the capital of France?")
print(list)
2.2.3、向量化检索模块（基于 TF-IDF）
文件：tf_idf.py
# new
# =============================================================================
# 第三部分：向量化模块（基于 TF-IDF）
# =============================================================================
import re
from collections import defaultdict
from typing import Dict, List, Tuple
import numpy as np
# 创建简单向量检索类
class SimpleVectorizer:
"""
简单的 TF-IDF向量化器
功能说明：
1. 将文本转换为TF-IDF向量
2. 支持基于余弦相似度的向量检索
3. 不依赖任何外部框架（如 sentence-transformers）
核心概念：
- TF (Term Frequency): 词频，某个词在文档中出现的频率
- IDF (Inverse Document Frequency):log(总文档/出现文档数+1) 逆文档频率，衡量词的重要性，值越大，表示在多个文档出现次数少
- TF-IDF = TF * IDF，值越大表示词对文档越重要
工作原理：
1. 分词：将文本切分成单词
2. 构建词表：收集所有出现过的词语
3. 计算 IDF：log(总文档数 / 包含该词的文档数)
4. 向量化：每个文档表示为一个稀疏向量
5. 检索：计算查询向量与所有文档向量的余弦相似度
使用场景：
- 小规模文档检索（几千篇以内）
- 不需要深度语义理解的场景
- 作为baseline(基线)或与语义检索融合
"""
def __init__(self, max_vocab_size: int = 1000):
"""
初始化向量化器
参数：
max_vocab_size: 最大词表大小，只保留高频词
"""
self.vocab: Dict[str, int] = {} # 词表：词 -> 索引
self.idf: Dict[str, float] = {} # IDF 值：词 -> 逆文档频率
self.corpus_tfidf: List[np.ndarray] = [] # 语料库的TF-IDF向量
self.documents: List[str] = [] # 原始文档列表
self.max_vocab_size = max_vocab_size # 词表大小限制
def tokenize(self, text: str) -> List[str]:
"""
文本分词（Tokenizer）
处理步骤：
1. 转小写（统一大小写）
2. 去除标点符号（保留字母和数字）
3. 按空格分割
参数：
text: 输入文本
返回：
单词列表 I am a handsome boy-》['i','am','a','handsome','boy']
示例：
tokenize("Hello, world!") -> ["hello", "world"]
"""
# 转换为小写
text = text.lower()
# 使用正则表达式：保留字母、数字、空格，移除其他字符
# [^\w\s]，表示非单词字符和空格之外的所有字符
text = re.sub(r'[^\w\s]', '', text)
# 按空白字符分割并返回
return text.split()
# 创建向量模型
def build_vocab(self, documents: List[str]) -> None:
"""
构建词表并计算 IDF
处理流程：
1. 遍历所有文档，统计：
- term_freq: 每个词出现的总次数
- doc_freq: 包含该词的文档数
2. 根据 term_freq 选择高频词构建词表
3. 计算每个词的 IDF 值
参数：
documents: 文档列表
"""
# 统计信息
term_freq = defaultdict(int) # 词频（所有文档中出现的总次数）
doc_freq = defaultdict(int) # 文档频率（包含该词的文档数）
# 第一遍遍历：统计词频和文档频率
for doc in documents:
# 使用 set 去重，确保每个文档对 doc_freq 的贡献为 1
unique_tokens = set(self.tokenize(doc))
for token in unique_tokens:
doc_freq[token] += 1
term_freq[token] += 1
# 构建词表：按词频降序排序，取前 max_vocab_size 个词
sorted_terms = sorted(term_freq.items(), key=lambda x: x[1], reverse=True)
top_terms = sorted_terms[:self.max_vocab_size]
# 创建词表映射（词 -> 索引）
self.vocab = {word: idx for idx, (word, _) in enumerate(top_terms)}
# 计算 IDF（逆文档频率）
# IDF 公式：log((N + 1) / (df + 1)) + 1
# 其中 N 是文档总数，df 是包含该词的文档数
# +1 是为了平滑，避免除零
n_docs = len(documents)
for word, idx in self.vocab.items():
df = doc_freq.get(word, 1)
self.idf[word] = np.log((n_docs + 1) / (df + 1)) + 1
def tfidf_vector(self, text: str) -> np.ndarray:
"""
计算文本的 TF-IDF向量
计算步骤：
1. 分词
2. 计算每个词的词频 TF = 词出现次数 / 总词数
3. 从预先计算的 IDF 中获取逆文档频率
4. 计算 TF * IDF 作为向量分量
参数：
text: 输入文本
返回：
TF-IDF 向量（numpy 数组，长度为词表大小）
"""
# 分词
tokens = self.tokenize(text)
if not tokens:
return np.zeros(len(self.vocab))
# 统计词频
token_counts = defaultdict(int)
for token in tokens:
token_counts[token] += 1
# 初始化零向量
vector = np.zeros(len(self.vocab))
# 计算每个词的 TF-IDF 值
for token, count in token_counts.items():
if token in self.vocab:
# TF = 词频 / 文档总词数
tf = count / len(tokens)
# 获取 IDF
idf = self.idf.get(token, 1.0)
# TF-IDF 加权
vector[self.vocab[token]] = tf * idf
# 返回自定义的词向量
return vector
def fit(self, documents: List[str]) -> None:
"""
拟合语料库（建立索引）
参数：
documents: 要索引的文档列表
"""
# 给检索文档添加索引号，以便后续检索
print(f" 正在为 {len(documents)} 个文档建立索引...")
self.build_vocab(documents)
self.corpus_tfidf = [self.tfidf_vector(doc) for doc in documents]
self.documents = documents
print(f" 索引建立完成，词表大小: {len(self.vocab)}")
def cosine_similarity(self, vec1: np.ndarray, vec2: np.ndarray) -> float:
"""
计算两个向量的余弦相似度
余弦相似度公式：cos(θ) = (A·B) / (|A| * |B|)
性质：
- 值域 [-1, 1]，但在 TF-IDF 中通常为 [0, 1]
- 1 表示方向完全相同（最相似）
- 0 表示正交（不相关）
参数：
vec1: 向量1
vec2: 向量2
返回：
余弦相似度值
"""
# 计算点积
dot_product = np.dot(vec1, vec2)
# 计算 L2 范数
norm1 = np.linalg.norm(vec1)
norm2 = np.linalg.norm(vec2)
# 避免除零错误
if norm1 == 0 or norm2 == 0:
return 0.0
# 返回余弦相似度：输入向量与文档向量
return dot_product / (norm1 * norm2)
def search(self, query: str, top_k: int = 5) -> List[Tuple[str, float]]:
"""
向量检索：从索引中查找与查询最相似的文档
参数：
query: 查询字符串
top_k: 返回的结果数量
返回：
文档和相似度分数的列表，按相似度降序排列
示例：
results = vectorizer.search("machine learning", top_k=3)
# 返回: [("文档1内容", 0.85), ("文档2内容", 0.72), ...]
"""
# 1. 将查询转换为向量
query_vec = self.tfidf_vector(query)
# 2. 计算与所有文档的相似度
scores = []
for i, doc_vec in enumerate(self.corpus_tfidf):
score = self.cosine_similarity(query_vec, doc_vec)
scores.append((self.documents[i], score))
# 3. 按相似度降序排序
scores.sort(key=lambda x: x[1], reverse=True)
# 4. 返回前top_k个结果：[{"文档1":90},{"文档2":95}]
return scores[:top_k]
# 测试一下，向量检索,使用tf-idf算法来决定被检索文档的匹配度
if __name__ == '__main__':
# 创建示例文档集合,模拟10个文档
documents = [
"Machine learning is a subset of artificial intelligence",
"Deep learning uses neural networks with many layers",
"Natural language processing helps computers understand text",
"Computer vision enables machines to interpret images",
"Reinforcement learning trains agents through rewards and penalties",
"Neural networks are inspired by biological brain structures",
"Data science combines statistics programming and domain expertise",
"Supervised learning requires labeled training data",
"Unsupervised learning finds patterns in unlabeled data",
"Transfer learning applies knowledge from one task to another"
]
# 创建向量化器实例
vectorizer = SimpleVectorizer(max_vocab_size=100)
# 建立索引
vectorizer.fit(documents)
print("\n" + "=" * 60)
print("测试1: 查询 'machine learning'")
print("=" * 60)
results = vectorizer.search("machine learning", top_k=3)
for i, (doc, score) in enumerate(results, 1):
print(f"{i}. 相似度: {score:.4f}")
print(f" 文档: {doc[:80]}...")
print("\n" + "=" * 60)
print("测试2: 查询 'neural networks'")
print("=" * 60)
results = vectorizer.search("neural networks", top_k=3)
for i, (doc, score) in enumerate(results, 1):
print(f"{i}. 相似度: {score:.4f}")
print(f" 文档: {doc[:80]}...")
print("\n" + "=" * 60)
print("测试3: 查询 'deep learning'")
print("=" * 60)
results = vectorizer.search("deep learning", top_k=3)
for i, (doc, score) in enumerate(results, 1):
print(f"{i}. 相似度: {score:.4f}")
print(f" 文档: {doc[:80]}...")
print("\n" + "=" * 60)
print("测试4: 查询 'artificial intelligence'")
print("=" * 60)
results = vectorizer.search("artificial intelligence", top_k=3)
for i, (doc, score) in enumerate(results, 1):
print(f"{i}. 相似度: {score:.4f}")
print(f" 文档: {doc[:80]}...")
2.2.4、关键词检索模块（BM25 算法）
文件：mybm25.py
#new
# =============================================================================
# 第四部分：关键词检索模块（BM25 算法）
# =============================================================================
import re
from collections import defaultdict # 导入构建字典工具
from typing import List, Dict, Tuple
import numpy as np
from nltk.corpus.reader import documents
from rank_bm25 import BM25
class SimpleBM25:
"""
BM25 (Okapi BM25) 检索器
BM25 是搜索引擎中经典的相关性评分算法，比简单的词频匹配更精确
核心公式：
score(D, Q) = Σ IDF(q_i) * (TF(q_i, D) * (k1 + 1)) / (TF(q_i, D) + k1 * (1 - b + b * |D|/avgdl))
参数说明：
- k1: 控制词频饱和度的参数（通常取 1.2-2.0）
- b: 控制文档长度归一化的参数（通常取 0.75）
- IDF: 逆文档频率
- TF: 词频
- |D|: 文档长度
- avgdl: 平均文档长度
优势：
1. 考虑了词频饱和（词出现太多不会无限增加分数）
2. 考虑了文档长度归一化（短文档匹配关键词更宝贵）
3. 在信息检索任务中表现优异
"""
def __init__(self, k1: float = 1.5, b: float = 0.75):
"""
初始化 BM25 检索器
参数：
k1: 词频饱和度参数，值越大 TF 影响越大
b: 长度归一化参数，值越大文档长度影响越大
"""
self.k1 = k1 # TF（词频） 饱和度参数
self.b = b # 长度归一化参数
self.documents: List[str] = [] # 原始文档
self.doc_lens: List[int] = [] # 每个文档的长度（词数）
self.avg_doc_len: float = 0.0 # 平均文档长度
self.inverted_index: Dict[str, List[Tuple[int, int]]] = defaultdict(list) # 倒排索引
self.idf: Dict[str, float] = {} # IDIF 缓存
def tokenize(self, text: str) -> List[str]:
"""
文本分词（优化版 - 更好支持中文）
"""
# 去除标点符号，但保留中文字符
text = re.sub(r'[^\w\u4e00-\u9fff]', ' ', text)
# 按空格分割
tokens = text.split()
# 进一步拆分长中文词（超过6个字符的）
result = []
for token in tokens:
if len(token) > 6 and re.match(r'^[\u4e00-\u9fff]+$', token):
# 长中文词，拆分成2-3字的词组
for i in range(len(token) - 1):
result.append(token[i:i+2])
if i + 3 <= len(token):
result.append(token[i:i+3])
else:
result.append(token)
# 去重
return list(set(result))
def index(self, documents: List[str]) -> None:
"""
构建 BM25 倒排索引
处理流程：
1. 计算每个文档的长度
2. 计算平均文档长度
3. 构建倒排索引：词 -> [(文档索引, 词频), ...]
4. 预计算每个词的 IDF
倒排索引原理：
正向索引：文档1 -> [词A, 词B, ...]
倒排索引：词A -> [文档1, 文档3, ...]
使用倒排索引可以快速找到包含查询词的文档，
而无需遍历所有文档。
参数：
documents: 文档列表
"""
self.documents = documents
n_docs = len(documents)
# 1. 计算文档长度和统计词频
term_freq_per_doc = [] # 每个文档的词频统计
for idx, doc in enumerate(documents):
tokens = self.tokenize(doc)
doc_len = len(tokens)
self.doc_lens.append(doc_len)
# 统计本文档的词频
tf_dict = defaultdict(int)
for token in tokens:
tf_dict[token] += 1
term_freq_per_doc.append(tf_dict)
# 2. 计算平均文档长度
self.avg_doc_len = np.mean(self.doc_lens)
# 3. 构建倒排索引（统计每个词出现在哪些文档及其词频）
for doc_id, tf_dict in enumerate(term_freq_per_doc):
for token, tf in tf_dict.items():
self.inverted_index[token].append((doc_id, tf))
# 4. 计算每个词的 IDF（逆文档频率）
# BM25 使用的 IDF 公式：log((N - df + 0.5) / (df + 0.5) + 1)
# 这个公式可以避免负的 IDF 值
for token, postings in self.inverted_index.items():
df = len(set(postings)) # 包含该词的文档数
self.idf[token] = np.log((n_docs - df + 0.5) / (df + 0.5) + 1)
print(f" BM25索引建立完成，包含 {len(self.inverted_index)} 个独特词语, {n_docs} 文档")
def search(self, query: str, top_k: int = 5) -> List[Tuple[str, float]]:
"""
BM25 检索
参数：
query: 查询字符串
top_k: 返回结果数量
返回：
文档和 BM25 分数的列表
"""
# 1. 对查询进行分词和去重
query_tokens = set(self.tokenize(query))
# 2. 初始化分数数组
scores = np.zeros(len(self.documents))
# 3. 对查询中的每个词计算 BM25 分数
for token in query_tokens:
if token not in self.inverted_index:
continue # 词不在索引中，跳过
idf = self.idf[token] # 获取词的 IDF
# 遍历包含该词的文档
for doc_id, tf in self.inverted_index[token]:
# 获取文档长度
doc_len = self.doc_lens[doc_id]
# BM25 核心公式
# 分子 = tf * (k1 + 1)
# 分母 = tf + k1 * (1 - b + b * doc_len/avgdl)
# 分数 += idf * 分子 / 分母
numerator = tf * (self.k1 + 1)
denominator = tf + self.k1 * (1 - self.b + self.b * doc_len / self.avg_doc_len)
scores[doc_id] += idf * numerator / denominator
# 4. 提取前 top_k 个结果
top_indices = np.argsort(scores)[::-1][:top_k]
results = []
for idx in top_indices:
if scores[idx] > 0:
results.append((self.documents[idx], float(scores[idx])))
return results
## 测试一下
# documents 变量未定义
# rank_bm25.BM25 是一个基类，不能直接实例化，需要使用具体的子类如 BM25Okapi
# rank_bm25 需要传入分词后的文档列表（二维列表），而不是原始字符串列表
# 创建示例文档集合10个
documents = [
"Machine learning is a subset of artificial intelligence",
"Deep learning uses neural networks with many layers",
"Natural language processing helps computers understand text",
"Computer vision enables machines to interpret images",
"Reinforcement learning trains agents through rewards and penalties",
"Neural networks are inspired by biological brain structures",
"Data science combines statistics programming and domain expertise",
"Supervised learning requires labeled training data",
"Unsupervised learning finds patterns in unlabeled data",
"Transfer learning applies knowledge from one task to another"
]
print("=" * 60)
print("测试 SimpleBM25（自定义实现）")
print("=" * 60)
# 使用自定义的 SimpleBM25
bm25_retriever = SimpleBM25(k1=1.5, b=0.75)
bm25_retriever.index(documents)
query = "machine learning"
print(f"\n查询: '{query}'")
results = bm25_retriever.search(query, top_k=3)
for i, (doc, score) in enumerate(results, 1):
print(f"{i}. 分数: {score:.4f}")
print(f" 文档: {doc[:80]}...")
query2 = "neural networks"
print(f"\n查询: '{query2}'")
results2 = bm25_retriever.search(query2, top_k=3)
for i, (doc, score) in enumerate(results2, 1):
print(f"{i}. 分数: {score:.4f}")
print(f" 文档: {doc[:80]}...")
print("\n" + "=" * 60)
print("测试 rank_bm25 库")
print("=" * 60)
# 使用 rank_bm25 库
from rank_bm25 import BM25Okapi
# 需要对文档进行分词
tokenized_docs = [doc.lower().split() for doc in documents]
bm25_okapi = BM25Okapi(tokenized_docs)
query_tokens = query.lower().split()
scores = bm25_okapi.get_scores(query_tokens)
print(f"\n查询: '{query}'")
for i, idx in enumerate(scores.argsort()[::-1][:3]):
if scores[idx] > 0:
print(f"{i + 1}. 分数: {scores[idx]:.4f}")
print(f" 文档: {documents[idx][:80]}...")
2.2.5、重排序模块（Re-Ranking）
文件：rerank.py
# =============================================================================
# 第五部分：重排序模块
# =============================================================================
import re
from dataclasses import dataclass, field
from typing import List
@dataclass
class Document:
"""
文档对象
"""
content: str = "" # 文档内容
score: float = 0.0 # 相关性分数
metadata: dict = field(default_factory=dict) # 元数据
def __repr__(self):
return f"Document(content={self.content[:50]}..., score={self.score:.4f})"
class ReRanker:
"""
重排序器（支持中文分词）
功能说明：
1. 对检索结果进行二次排序
2. 基于词语重叠和短语匹配计算相关性
3. 支持中英文混合文本
工作原理：
1. 分词：对查询和文档进行智能分词
2. 计算词语重叠率
3. 检查短语匹配度
4. 综合评分并重新排序
"""
def __init__(self, k1: float = 1.5, b: float = 0.75):
"""
初始化重排序器
参数：
k1: 词频饱和度参数（保留用于未来扩展）
b: 长度归一化参数（保留用于未来扩展）
"""
# 停用词（常见无意义的词）
self.stopwords = set([
'的', '了', '是', '在', '我', '有', '和', '就', '不', '人',
'都', '一', '一个', '上', '也', '很', '到', '说', '要', '去',
'你', '会', '着', '没有', '看', '好', '自己', '这', '那',
'什么', '如何', '为什么', '怎么', '哪个', '哪些', '请',
'介绍', '一下', '它', '他', '她', '我们', '他们', '这个',
'那个', '可以', '可能', '应该', '如果', '但是', '因为',
'所以', '虽然', '然后', '接着', '于是', '还', '又', '再'
])
def tokenize_chinese(self, text: str) -> List[str]:
"""
简单的中文分词
策略：
1. 按标点符号分割成短语
2. 提取有意义的词语（长度>=2，不是纯数字，不是停用词）
参数：
text: 输入文本
返回：
分词后的列表
示例：
tokenize_chinese("周杰伦的七里香很好听")
-> ["周杰伦", "七里香", "好听"]
"""
# 去除标点符号和空白字符（使用原始字符串避免转义警告）
text = re.sub(r'[，。！？；：""''（）【】《》、\s]+', ' ', text)
# 按空格分割
words = text.split()
# 过滤和切分
tokens = []
for word in words:
# 跳过纯数字和过短的词
if word.isdigit() or len(word) < 2:
continue
# 过滤停用词
if word not in self.stopwords:
tokens.append(word)
return tokens
def tokenize_simple(self, text: str) -> List[str]:
"""
简单分词：按字符拆分（用于短文本匹配）
策略：
1. 只保留中文字符
2. 生成2-4字的滑动窗口词组
参数：
text: 输入文本
返回：
分词后的列表
"""
# 去除标点符号和空格，只保留中文
text = re.sub(r'[^\u4e00-\u9fff]', '', text)
# 按字符拆分，但合并成2-4个字符的词组
tokens = []
length = len(text)
# 加入完整文本
if text:
tokens.append(text)
# 加入2-4字的子串
for i in range(length):
for j in range(2, min(5, length - i + 1)):
substr = text[i:i + j]
if len(substr) >= 2 and substr not in self.stopwords:
tokens.append(substr)
return list(set(tokens))
def compute_token_overlap(self, query: str, document: str) -> float:
"""
计算查询与文档的词语重叠率（支持中文）
参数：
query: 查询字符串
document: 文档内容
返回：
重叠率（0.0-1.0）
"""
# 使用中文分词
query_tokens = self.tokenize_chinese(query)
doc_tokens = self.tokenize_chinese(document)
if not query_tokens:
# 如果分词结果为空，尝试简单分词
query_tokens = self.tokenize_simple(query)
doc_tokens = self.tokenize_simple(document)
if not query_tokens:
return 0.0
# 转换为集合
query_set = set(query_tokens)
doc_set = set(doc_tokens)
# 计算交集大小
overlap = len(query_set & doc_set)
# 重叠率 = 交集大小 / 查询词数
return overlap / len(query_set)
def check_phrase_match(self, query: str, document: str) -> float:
"""
检查查询中的关键词是否在文档中出现
参数：
query: 查询字符串
document: 文档内容
返回：
匹配分数（0.0-1.0）
"""
# 提取查询中的关键词（去除标点）
query_clean = re.sub(r'[？?！!。，,、；;：:]', '', query)
# 如果查询完整出现在文档中
if query_clean in document:
return 1.0
# 提取文档中的关键词
doc_keywords = self.tokenize_chinese(document)
# 将查询分成多个部分检查
parts = query_clean.split()
for part in parts:
if len(part) >= 2 and part in document:
return 0.8
# 检查查询中的关键词
query_keywords = self.tokenize_chinese(query_clean)
for kw in query_keywords:
if kw in doc_keywords or kw in document:
return 0.6
return 0.0
def compute_relevance(self, query: str, document: str) -> float:
"""
综合计算查询与文档的相关性分数
参数：
query: 查询字符串
document: 文档内容
返回：
相关性分数（0.0-1.0）
"""
# 基础词语重叠分
base_score = self.compute_token_overlap(query, document)
# 短语匹配加分
phrase_bonus = self.check_phrase_match(query, document)
# 加权综合（短语匹配权重更高）
if phrase_bonus >= 1.0:
return max(base_score, 0.8)
elif phrase_bonus > 0.5:
return base_score * 0.3 + phrase_bonus * 0.7 # 增加短语匹配权重
else:
return base_score * 0.5 + phrase_bonus * 0.5 # 平衡权重
def rerank(self, query: str, documents: List[Document]) -> List[Document]:
"""
对检索结果进行重排序
参数：
query: 查询字符串
documents: 待重排序的文档列表
返回：
重排序后的文档列表
"""
# 计算每个文档的新相关性分数
for doc in documents:
relevance = self.compute_relevance(query, doc.content)
# 综合原始分数和相关性分数（70% 新分数 + 30% 原始分数）
doc.score = relevance * 0.7 + doc.score * 0.3
# 按新分数重新排序
documents.sort(key=lambda x: x.score, reverse=True)
return documents
# ===== 测试代码 =====
if __name__ == "__main__":
print("=" * 60)
print("测试重排序器")
print("=" * 60)
# 创建重排序器实例
reranker = ReRanker()
# 测试查询
query = "如何评价周杰伦的《七里香》？"
# 创建测试文档
documents = [
Document(content="周杰伦的《七里香》是一张非常经典的专辑，融合了多种音乐风格", score=0.8),
Document(content="机器学习是人工智能的一个分支，涉及算法和统计模型", score=0.6),
Document(content="深度学习使用多层神经网络来处理复杂的数据模式", score=0.5),
Document(content="自然语言处理技术可以帮助计算机理解人类语言", score=0.4),
Document(content="周杰伦专辑《七里香》的专辑简介", score=0.3),
Document(content="周杰伦专辑《七里香》的专辑简介", score=0.2),
Document(content="周杰伦专辑《七里香》的专辑简介", score=0.1)
]
# 输出重排后的文档
list=reranker.rerank(query,documents)
print()
# 取前Top-3
print(list[:3])
2.2.6、上下文压缩模块
文件：myprompress.py
# =============================================================================
# 第六部分：上下文压缩模块
# =============================================================================
import re
from dataclasses import dataclass, field
from typing import List
@dataclass
class Document:
"""
文档对象
"""
content: str = "" # 文档内容
score: float = 0.0 # 相关性分数
metadata: dict = field(default_factory=dict) # 元数据
def __repr__(self):
return f"Document(content={self.content[:50]}..., score={self.score:.4f})"
class ContextCompressor:
"""
上下文压缩器
为什么需要压缩？
1. LLM 的上下文窗口有限（如 4K、8K、128K tokens）
2. 检索到的文档可能很长，包含大量无关信息
3. 压缩后的上下文更精准，生成质量更高
4. 节省 token 成本
压缩策略：
1. 分句：将文档拆分成句子
2. 评分：评估每个句子与查询的相关性
3. 选择：只保留最相关的几个句子
4. 重组：按原文顺序重新组织
"""
def __init__(self, max_sentences: int = 3):
"""
初始化压缩器
参数：
max_sentences: 每个文档保留的最大句子数
"""
self.max_sentences = max_sentences
def split_sentences(self, text: str) -> List[str]:
"""
将文本拆分成句子
支持中英文标点符号作为句子分隔符
参数：
text: 输入文本
返回：
句子列表
"""
# 使用正则表达式按句子结束符分割（支持中英文）
sentences = re.split(r'[.!?。！？]+', text)
# 过滤掉空字符串和太短的句子
return [s.strip() for s in sentences if len(s.strip()) > 10]
def compute_sentence_relevance(self, query: str, sentence: str) -> float:
"""
计算句子与查询的相关性
使用 Token 重叠率作为相关性指标
参数：
query: 查询字符串
sentence: 句子字符串
返回：
相关性分数（0-1）
"""
# 简单分词：按空格或字符分割
query_tokens = set(query.lower().split())
sent_tokens = set(sentence.lower().split())
# 如果没有空格（中文），按字符分割
if not query_tokens and query:
query_tokens = set(query)
if not sent_tokens and sentence:
sent_tokens = set(sentence)
if not query_tokens:
return 0.0
overlap = len(query_tokens & sent_tokens)
return overlap / len(query_tokens)
def compress(self, query: str, documents: List[Document]) -> List[Document]:
"""
压缩文档列表
处理流程：
1. 对每个文档，拆分成句子
2. 计算每个句子的相关性分数
3. 取出分数最高的几个句子
4. 用这些句子重构文档
参数：
query: 查询字符串
documents: 输入的文档列表
返回：
压缩后的文档列表
"""
compressed_docs = []
for doc in documents:
# 拆分文档为句子
sentences = self.split_sentences(doc.content)
if not sentences:
# 如果拆不出句子，保持原样
compressed_docs.append(doc)
continue
# 计算每个句子的相关性
scored_sentences = []
for sent in sentences:
score = self.compute_sentence_relevance(query, sent)
scored_sentences.append((sent, score))
# 按相关性分数降序排序
scored_sentences.sort(key=lambda x: x[1], reverse=True)
# 取 top-k 最相关的句子
top_sentences = [s for s, _ in scored_sentences[:self.max_sentences]]
# 为了保持可读性，按原文顺序重新组织
# 这里简单处理：直接拼接
compressed_content = " ".join(top_sentences)
# 创建压缩后的文档
compressed_doc = Document(
content=compressed_content,
metadata=doc.metadata,
score=doc.score
)
compressed_docs.append(compressed_doc)
return compressed_docs
# 测试一下
if __name__ == "__main__":
print("=" * 60)
print("测试上下文压缩器")
print("=" * 60)
# 创建测试文档'test.txt'
doc = Document(
content="Machine learning is a subset of artificial intelligence. It focuses on building systems that learn from data. Deep learning is a type of machine learning. Neural networks are the foundation of deep learning.",
score=0.9,
metadata={"source": "test.txt"}
)
# 创建压缩器实例
comp = ContextCompressor(max_sentences=2)
# 测试查询
query = "deep learning neural networks"
print(f"\n原始文档:")
print(f" 内容长度: {len(doc.content)} 字符")
print(f" 内容: {doc.content}")
# 压缩文档
compressed_docs = comp.compress(query, [doc])
print(f"\n压缩后文档:")
for i, compressed_doc in enumerate(compressed_docs, 1):
print(f" 内容长度: {len(compressed_doc.content)} 字符")
print(f" 内容: {compressed_doc.content}")
print(f" 分数: {compressed_doc.score}")
print(f" 元数据: {compressed_doc.metadata}")
print("\n" + "=" * 60)
print("测试中文文档压缩")
print("=" * 60)
# 测试中文文档:'chinese_test.txt'
chinese_doc = Document(
content="机器学习是人工智能的一个分支。它专注于构建从数据中学习的系统。深度学习是机器学习的一种类型。神经网络是深度学习的基础。自然语言处理是另一个重要领域。",
score=0.85,
metadata={"source": "chinese_test.txt"}
)
chinese_query = "深度学习 神经网络"
print(f"\n原始文档:")
print(f" 内容长度: {len(chinese_doc.content)} 字符")
print(f" 内容: {chinese_doc.content}")
compressed_chinese_docs = comp.compress(chinese_query, [chinese_doc])
print(f"\n压缩后文档:")
for i, compressed_doc in enumerate(compressed_chinese_docs, 1):
print(f" 内容长度: {len(compressed_doc.content)} 字符")
print(f" 内容: {compressed_doc.content}")
print(f" 分数: {compressed_doc.score}")
print(f" 元数据: {compressed_doc.metadata}")
2.2.7、融合模块（Fusion）
文件：myfusion.py
# new
# =============================================================================
# 第七部分：融合模块（Fusion）
# =============================================================================
from collections import defaultdict
from typing import List, Optional
from qdrant_client.grpc import Document
class FusionModule:
"""
多路检索结果融合模块
为什么需要融合？
- 不同的检索器有不同的优势：
* 向量检索：捕获语义相似性
* BM25：精确匹配关键词
- 融合可以取长补短，获得更好的召回效果
支持的融合策略：
1. RRF (Reciprocal Rank Fusion): 最常用的融合方法（倒数排序融合）
2. Weighted Fusion: 加权融合，可以为不同检索器设置不同权重
3. Concatenation: 简单拼接 + 去重
"""
def __init__(self, method: str = "rrf"):
"""
初始化融合器
参数：
method: 融合方法，可选 "rrf", "weighted", "concat"
"""
self.method = method # 融合策略选择
def reciprocal_rank_fusion(self,
results_list: List[List[Document]],
k: int = 60) -> List[Document]:
"""
RRF (Reciprocal Rank Fusion) 算法
这是最常用的融合算法，在 TREC 等评测中表现优异
公式：
RRF_score(d) = Σ 1 / (k + rank(d))
其中：
- rank(d) 是文档 d 在某个检索结果中的排名（从 0 开始）
- k 是常数（通常取 60，用于平滑）
特点：
- 不需要归一化分数（不同检索器的分数量纲可能不同）
- 只使用排名信息，鲁棒性好
- 高排名（小 rank）的文档贡献更大
参数：
results_list: 多路检索结果列表，每路结果都是 Document 列表
k: RRF 常数，通常取 60
返回：
融合后的文档列表（按 RRF 分数降序排列）
"""
rrf_scores = defaultdict(float) # 存储每个文档的 RRF 分数
doc_map = {} # 文档内容到 Document 对象的映射
# 遍历每一路检索结果
for rank_list in results_list:
# 只考虑前 20 个结果（排名靠后的贡献很小）
for rank, doc in enumerate(rank_list[:20]):
doc_key = doc.content # 使用内容作为唯一标识
# RRF 贡献：1 / (k + rank)，rank 从 0 开始，所以实际用 rank
rrf_scores[doc_key] += 1.0 / (k + rank + 1) # +1 使 rank 从 1 开始
# 保存文档对象（如果还没有保存过）
if doc_key not in doc_map:
doc_map[doc_key] = doc
# 构建融合后的结果列表
merged_results = []
for doc_key, score in sorted(rrf_scores.items(), key=lambda x: x[1], reverse=True):
doc = doc_map[doc_key]
doc.score = score # 更新分数
merged_results.append(doc)
return merged_results
def weighted_fusion(self,
results_list: List[List[Document]],
weights: List[float]) -> List[Document]:
"""
加权融合
原理：为不同的检索器设置不同的权重，分数 = Σ weight_i * normalized_score_i
适用场景：
- 知道哪个检索器更可靠
- 需要突出某个检索器的结果
参数：
results_list: 多路检索结果列表
weights: 权重列表，长度必须与 results_list 相同，总和应为 1
返回：
融合后的文档列表
"""
weighted_scores = defaultdict(float)
doc_map = {}
# 遍历每一路检索结果
for weight, rank_list in zip(weights, results_list):
for rank, doc in enumerate(rank_list[:10]):
doc_key = doc.content
# 分数 = 权重 * (1 / (rank + 1))，实现排名衰减
score = weight * (1.0 / (rank + 1))
weighted_scores[doc_key] += score
if doc_key not in doc_map:
doc_map[doc_key] = doc
# 排序并返回
merged_results = []
for doc_key, score in sorted(weighted_scores.items(), key=lambda x: x[1], reverse=True):
doc = doc_map[doc_key]
doc.score = score
merged_results.append(doc)
return merged_results
def simple_concat(self, results_list: List[List[Document]]) -> List[Document]:
"""
简单拼接去重
原理：
1. 按顺序拼接各路结果
2. 去除重复文档（基于内容）
3. 保持原有的相对顺序
参数：
results_list: 多路检索结果列表
返回：
去重后的文档列表
"""
seen = set()
merged = []
for rank_list in results_list:
for doc in rank_list:
if doc.content not in seen:
seen.add(doc.content)
merged.append(doc)
return merged
def fuse(self,
results_list: List[List[Document]],
weights: Optional[List[float]] = None) -> List[Document]:
"""
融合入口方法
参数：
results_list: 多路检索结果列表
weights: 权重列表（仅 weighted 方法需要）
返回：
融合后的文档列表
"""
if self.method == "rrf":
return self.reciprocal_rank_fusion(results_list)
elif self.method == "weighted":
if weights is None:
weights = [1.0 / len(results_list)] * len(results_list)
return self.weighted_fusion(results_list, weights)
else: # concat
return self.simple_concat(results_list)
# 测试一下=================
if __name__ == "__main__":
print("=" * 60)
print("测试融合模块")
print("=" * 60)
# 创建自定义 Document 类用于测试
from dataclasses import dataclass, field
@dataclass
class SimpleDocument:
content: str = ""
score: float = 0.0
metadata: dict = field(default_factory=dict)
def __repr__(self):
return f"Document(content={self.content[:50]}..., score={self.score:.4f})"
# 模拟两路检索结果
vector_results = [
SimpleDocument(content="机器学习是人工智能的一个分支", score=0.9),
SimpleDocument(content="深度学习使用神经网络", score=0.8),
SimpleDocument(content="自然语言处理技术", score=0.7),
SimpleDocument(content="计算机视觉应用", score=0.6),
SimpleDocument(content="强化学习方法", score=0.5)
]
bm25_results = [
SimpleDocument(content="机器学习算法与应用", score=0.85),
SimpleDocument(content="深度学习框架介绍", score=0.75),
SimpleDocument(content="数据挖掘技术", score=0.65),
SimpleDocument(content="神经网络基础", score=0.55),
SimpleDocument(content="人工智能发展史", score=0.45)
]
print("\n向量检索结果:")
for i, doc in enumerate(vector_results, 1):
print(f" {i}. {doc.content} (分数: {doc.score:.2f})")
print("\nBM25检索结果:")
for i, doc in enumerate(bm25_results, 1):
print(f" {i}. {doc.content} (分数: {doc.score:.2f})")
# 测试 RRF 融合
print("\n" + "=" * 60)
print("测试1: RRF 融合")
print("=" * 60)
fusion_module = FusionModule(method="rrf")
rrf_results = fusion_module.fuse([vector_results, bm25_results])
print(f"\nRRF 融合结果（前5个）:")
for i, doc in enumerate(rrf_results[:5], 1):
print(f" {i}. {doc.content} (RRF分数: {doc.score:.4f})")
# 测试加权融合
print("\n" + "=" * 60)
print("测试2: 加权融合")
print("=" * 60)
fusion_module_weighted = FusionModule(method="weighted")
weighted_results = fusion_module_weighted.fuse(
[vector_results, bm25_results],
weights=[0.6, 0.4] # 向量检索权重更高
)
print(f"\n加权融合结果（前5个，权重: 向量0.6, BM25 0.4）:")
for i, doc in enumerate(weighted_results[:5], 1):
print(f" {i}. {doc.content} (加权分数: {doc.score:.4f})")
# 测试简单拼接
print("\n" + "=" * 60)
print("测试3: 简单拼接去重")
print("=" * 60)
fusion_module_concat = FusionModule(method="concat")
concat_results = fusion_module_concat.fuse([vector_results, bm25_results])
print(f"\n拼接去重结果（共{len(concat_results)}个文档）:")
for i, doc in enumerate(concat_results[:5], 1):
print(f" {i}. {doc.content}")
print("\n" + "=" * 60)
print("融合测试完成！")
print("=" * 60)
2.2.8、生成器模块（Generator）
文件：my_generator.py
# new
# =============================================================================
# 第八部分：生成器模块（Generator）
# =============================================================================
from typing import List
from qdrant_client.grpc import Document
class SimpleGenerator:
"""
生成器模块（模拟大语言模型）
在真实的 RAG 系统中，这里应该调用：
- OpenAI GPT API
- Anthropic Claude API
- 本地模型（如 Llama、ChatGLM）
- 或其他 LLM 服务
这里实现一个模拟生成器，包含：
1. Prompt 构建
2. 模拟的知识库查询
3. 基于上下文的简单答案生成
真实实现示例（伪代码）：
def generate(self, query, documents):
context = "\n".join([doc.content for doc in documents[:3]])
prompt = f"Context: {context}\n\nQuestion: {query}\n\nAnswer:"
response = openai.ChatCompletion.create(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": prompt}]
)
return response.choices[0].message.content
"""
def __init__(self, use_mock: bool = True):
"""
初始化生成器
参数：
use_mock: 是否使用模拟模式（True 用于测试，False 需要接入真实 LLM）
"""
self.use_mock = use_mock
# 模拟的知识库（仅用于演示）
self.mock_knowledge = {
"python": "Python 是一种高级、解释型的编程语言，以其简洁和可读性而闻名。",
"rag": "RAG（检索增强生成）将信息检索与大语言模型相结合，以生成更准确、信息更丰富的回答。",
"machine learning": "机器学习是人工智能的一个子集，使系统能够从数据中学习而无需显式编程。",
"chatgpt": "ChatGPT 是 OpenAI 开发的对话式人工智能模型，基于 GPT 架构。",
"向量检索": "向量检索是一种使用向量表示文本，并通过计算向量相似度来查找相关文档的技术。",
"bm25": "BM25 是一种经典的文本检索算法，广泛应用于搜索引擎中计算文档相关性。",
"什么是": "这是一个查询开头词，用于询问定义或解释。"
}
def build_prompt(self, query: str, contexts: List[str]) -> str:
"""
构建发送给 LLM 的提示词
一个好的 prompt 应该包含：
1. 系统指令（告诉模型如何回答）
2. 上下文信息
3. 用户问题
4. 输出格式要求
参数：
query: 用户查询
contexts: 上下文文档列表
返回：
完整的 prompt 字符串
"""
# 构建上下文部分
context_text = ""
for i, ctx in enumerate(contexts):
context_text += f"\n[Document {i+1}]: {ctx}\n"
# 完整的 prompt 模板
prompt = f"""你是一个乐于助人的助手，根据提供的上下文回答问题。
请只使用上下文中的信息来回答用户的问题。
如果上下文信息不足以回答问题，请说"我没有足够的信息来回答这个问题"。
上下文：
{context_text}
问题：{query}
请基于上述上下文回答问题："""
return prompt
def mock_generate(self, query: str, contexts: List[str]) -> str:
"""
模拟回答生成
实际应用中，这里应该是 LLM API 调用
这里实现了简单的关键字匹配 + 上下文提取
参数：
query: 查询
contexts: 上下文
返回：
生成的答案
"""
query_lower = query.lower()
# 1. 尝试从模拟知识库中匹配
for keyword, answer in self.mock_knowledge.items():
if keyword in query_lower:
# 如果有上下文，结合上下文给出更完整的回答
if contexts:
context_hint = "\n\nBased on the provided context, " + answer
return f"根据提供的上下文: {answer}"
return answer
# 2. 基于上下文生成回答
if contexts:
# 合并上下文
combined_context = " ".join(contexts)
# 尝试在上下文中找到与查询相关的部分
# 简单策略：查找包含查询关键词的句子
sentences = combined_context.split('.')
relevant_sentences = []
for sent in sentences:
# 检查句子是否包含查询中的关键词
if any(word in sent.lower() for word in query_lower.split()[:2]):
relevant_sentences.append(sent.strip())
if relevant_sentences:
# 取最相关的前2个句子
answer = ". ".join(relevant_sentences[:2])
if answer:
return f"根据所提供的信息: {answer}."
# 如果没有找到相关句子，返回上下文摘要
preview = combined_context[:200]
return f"根据上下文信息：{preview}......"
# 3. 没有任何信息可用
return f"我没有足够的信息来回答'{query}'。请提供更多上下文或检查知识库。"
def generate(self, query: str, documents: List[Document]) -> str:
"""
生成回答（主入口）
参数：
query: 用户查询
documents: 检索到的文档列表（已排序）
返回：
生成的答案文本
"""
# 只使用前3个最相关的文档作为上下文
context_docs = documents[:3]
contexts = [doc.content for doc in context_docs]
if self.use_mock:
# 模拟模式
return self.mock_generate(query, contexts)
else:
# 真实模式：这里接入真实的 LLM API
# 示例代码（需要安装 openai 包）
"""
import openai
prompt = self.build_prompt(query, contexts)
response = openai.ChatCompletion.create(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": prompt}],
temperature=0.7
)
return response.choices[0].message.content
"""
return "[Real LLM integration would go here]"
# ====测试一下===================
if __name__ == "__main__":
from dataclasses import dataclass, field
# 定义简单的 Document 类用于测试
@dataclass
class SimpleDocument:
content: str = ""
score: float = 0.0
metadata: dict = field(default_factory=dict)
def __repr__(self):
return f"Document(content={self.content[:50]}..., score={self.score:.4f})"
print("=" * 60)
print("测试生成器模块")
print("=" * 60)
# 创建生成器实例
generator = SimpleGenerator(use_mock=True)
# 测试1: 英文查询
print("\n测试1: 英文查询")
print("-" * 40)
query1 = "What is Python?"
documents1 = [
SimpleDocument(content="Python is a high-level programming language known for its simplicity.", score=0.9),
SimpleDocument(content="Python supports multiple programming paradigms.", score=0.8),
SimpleDocument(content="Machine learning libraries are often written in Python.", score=0.7)
]
answer1 = generator.generate(query1, documents1)
print(f"查询: {query1}")
print(f"回答: {answer1}\n")
# 测试2: 中文查询
print("测试2: 中文查询")
print("-" * 40)
query2 = "什么是机器学习？"
documents2 = [
SimpleDocument(content="机器学习是人工智能的一个子集，使系统能够从数据中学习而无需显式编程。", score=0.95),
SimpleDocument(content="深度学习是机器学习的一种特殊形式，使用多层神经网络。", score=0.85),
SimpleDocument(content="监督学习需要标注数据，无监督学习不需要标注数据。", score=0.75)
]
answer2 = generator.generate(query2, documents2)
print(f"查询: {query2}")
print(f"回答: {answer2}\n")
# 测试3: RAG相关查询
print("测试3: RAG相关查询")
print("-" * 40)
query3 = "什么是RAG技术？"
documents3 = [
SimpleDocument(content="RAG（检索增强生成）将信息检索与大语言模型相结合，以生成更准确、信息更丰富的回答。",
score=0.92),
SimpleDocument(content="向量检索是一种使用向量表示文本的技术。", score=0.7),
SimpleDocument(content="BM25算法常用于关键词检索。", score=0.6)
]
answer3 = generator.generate(query3, documents3)
print(f"查询: {query3}")
print(f"回答: {answer3}\n")
# 测试4: 无足够信息的查询
print("测试4: 无足够信息的查询")
print("-" * 40)
query4 = "量子计算机的工作原理是什么？"
documents4 = [
SimpleDocument(content="这是一篇关于机器学习的文章。", score=0.3),
SimpleDocument(content="Python是一种编程语言。", score=0.2)
]
answer4 = generator.generate(query4, documents4)
print(f"查询: {query4}")
print(f"回答: {answer4}\n")
print("=" * 60)
print("生成器测试完成！")
print("=" * 60)
2.2.9、Modular RAG 主类与整体测试
文件：demo_main.py
# new
# =============================================================================
# 第九部分：Modular RAG 主类
# =============================================================================
from typing import Optional, Dict, List, Any
from dataclasses import dataclass, field
#==================导入自定义的模块===============
from query import query_rewriter as qr
from retrival import mybm25, tf_idf
from post_processing import mypromress, rerank
from fusion_generator import my_generator, myfusion
# 定义自定义 Document 类（替代 qdrant_client.grpc.Document）
@dataclass
class Document:
"""
文档对象
"""
content: str = "" # 文档内容
score: float = 0.0 # 相关性分数
metadata: dict = field(default_factory=dict) # 元数据
def __repr__(self):
return f"Document(content={self.content[:50]}..., score={self.score:.4f})"
# 调用所有模块的类
class ModularRAG:
"""
Modular RAG（模块化检索增强生成）主类
这是一个完整的 RAG 系统实现，所有组件都是可插拔、可配置的
系统架构：
┌─────────────────────────────────────────────────────────────────┐
│ 用户输入查询 │
└─────────────────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────┐
│ Module 1: Query Rewriter（查询改写） │
│ - 同义词扩展 / HyDE / 多查询生成 │
└─────────────────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────┐
│ Module 2: Multi-Retriever（多路检索） │
│ - 向量检索（TF-IDF） │
│ - 关键词检索（BM25） │
└─────────────────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────┐
│ Module 3: Fusion（结果融合） │
│ - RRF（倒排序融合） / 加权融合 / 拼接 │
└─────────────────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────┐
│ Module 4: ReRanker（重排序） │
│ - 精细化相关性计算 │
└─────────────────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────┐
│ Module 5: Context Compressor（上下文压缩） │
│ - 提取最相关的句子 │
└─────────────────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────┐
│ Module 6: Generator（生成器） │
│ - 构建 Prompt + LLM 生成答案 │
└─────────────────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────┐
│ 最终答案 │
└─────────────────────────────────────────────────────────────────┘
使用示例：
rag = ModularRAG()
rag.index_documents(documents)
result = rag.query("什么是RAG?")
print(result["answer"])
"""
#### 初始化配置
def __init__(self, config: Optional[Dict] = None):
"""
初始化 Modular RAG 系统
参数：
config: 配置字典，可以自定义各个组件的行为
{
"fusion_method": "rrf", # rrf, weighted, concat
"max_sentences": 3, # 每个文档保留的最大句子数
"enable_rewrite": True, # 是否启用查询改写
"enable_rerank": True, # 是否启用重排序
"enable_fusion": True, # 是否启用融合
"enable_compression": True, # 是否启用压缩
"rewrite_method": "multi" # expand, hyde, multi
}
"""
self.config = config or {}
#### 初始化各个模块
print("初始化Modular RAG的每个模块...")
# 1）、索引向量
self.vectorizer = tf_idf.SimpleVectorizer(max_vocab_size=1000)
# 2）、向量检索
self.bm25 = mybm25.SimpleBM25(k1=1.5, b=0.75)
# 3）、重排序
self.reranker = rerank.ReRanker()
# 4）、融合
self.fusion = myfusion.FusionModule(method=self.config.get("fusion_method", "rrf"))
# 5）、查询改写
self.rewriter = qr.QueryRewriter()
# 6）、上下文压缩
self.compressor = mypromress.ContextCompressor(
max_sentences=self.config.get("max_sentences", 3)
)
# 7）、生成器
self.generator = my_generator.SimpleGenerator(use_mock=True)
self.is_indexed = False # 文档索引状态，默认为未索引
print("初始化完成 !")
# 索引文档到所有检索器
def index_documents(self, documents: List[str]) -> None:
"""
索引文档到所有检索器
参数：
documents: 要索引的文档列表
"""
print(f"\n正在为 {len(documents)} 个文档建立索引...")
self.vectorizer.fit(documents)
self.bm25.index(documents)
self.is_indexed = True
print("索引完成!\n")
# 使用向量检索
def search_vector(self, query: str, top_k: int = 5) -> List[Document]:
"""
使用向量检索
参数：
query: 查询字符串
top_k: 返回数量
返回：
Document 对象列表
"""
results = self.vectorizer.search(query, top_k)
return [Document(content=doc, score=score) for doc, score in results]
# 使用 BM25 关键词检索
def search_keyword(self, query: str, top_k: int = 5) -> List[Document]:
"""
使用 BM25 关键词检索
参数：
query: 查询字符串
top_k: 返回数量
返回：
Document 对象列表
"""
results = self.bm25.search(query, top_k)
return [Document(content=doc, score=score) for doc, score in results]
# 执行多路检索
def multi_retrieve(self, query: str, top_k: int = 5) -> List[List[Document]]:
"""
执行多路检索
返回：
[向量检索结果, BM25检索结果]
"""
vector_results = self.search_vector(query, top_k)
keyword_results = self.search_keyword(query, top_k)
return [vector_results, keyword_results]
# 执执完整的RAG查询
def query(self,
query: str,
enable_rewrite: bool = None,
enable_fusion: bool = None,
enable_rerank: bool = None,
enable_compression: bool = None,
rewrite_method: str = None,
top_k: int = 5,
verbose: bool = True) -> Dict[str, Any]:
"""
执行完整的 RAG 查询流程
参数：
query: 用户查询
enable_rewrite: 是否启用查询改写
enable_fusion: 是否启用多路融合
enable_rerank: 是否启用重排序
enable_compression: 是否启用上下文压缩
rewrite_method: 改写方法 ("expand", "hyde", "multi")
top_k: 每路检索返回的文档数
verbose: 是否打印详细过程
返回：
包含查询结果的字典：
{
"query": 原始查询,
"rewritten_queries": 改写后的查询列表,
"retrieved_count": 检索到的文档总数,
"fused_count": 融合后的文档数,
"final_documents": 最终的文档列表,
"answer": 生成的答案
}
"""
if not self.is_indexed:
return {"错误": "请先调用 index_documents() 方法"}
# 使用配置或默认值
enable_rewrite = enable_rewrite if enable_rewrite is not None else self.config.get("enable_rewrite", True)
enable_fusion = enable_fusion if enable_fusion is not None else self.config.get("enable_fusion", True)
enable_rerank = enable_rerank if enable_rerank is not None else self.config.get("enable_rerank", True)
enable_compression = enable_compression if enable_compression is not None else self.config.get(
"enable_compression", True)
rewrite_method = rewrite_method or self.config.get("rewrite_method", "multi")
if verbose:
print(f"\n{'=' * 60}")
print(f"Query: {query}")
print(f"{'=' * 60}\n")
# ============================================================
# Step 1: 查询改写
# ============================================================
if enable_rewrite:
queries = self.rewriter.rewrite(query, method=rewrite_method)
if verbose:
print(f"✏️ Step 1 - 查询改写:")
print(f" 原始查询: {query}")
print(f" 扩展查询(改写后的查询): {queries}")
else:
queries = [query]
if verbose:
print(f"✏️ Step 1 - 查询改写被禁用")
# ============================================================
# Step 2: 多路检索
# ============================================================
if verbose:
print(f"\n🔍 Step 2 - 多路检索:")
all_results = []
for q in queries:
vector_res = self.search_vector(q, top_k=top_k)
keyword_res = self.search_keyword(q, top_k=top_k)
all_results.extend([vector_res, keyword_res])
if verbose:
print(f" 查询 '{q[:30]}...': 向量检索={len(vector_res)}, bm25={len(keyword_res)}")
total_retrieved = sum(len(r) for r in all_results)
if verbose:
print(f" 原始结果总数: {total_retrieved}")
# ============================================================
# Step 3: 结果融合
# ============================================================
if enable_fusion and len(all_results) > 1:
fused_results = self.fusion.fuse(all_results)
if verbose:
print(f"\n🔄 Step 3 - 结果融合 ({self.fusion.method}):")
print(f" 融合后文档数: {len(fused_results)}")
else:
fused_results = all_results[0] if all_results else []
if verbose:
print(f"\n🔄 Step 3 - 融合: 已禁用")
# ============================================================
# Step 4: 重排序
# ============================================================
if enable_rerank:
reranked_results = self.reranker.rerank(query, fused_results)
if verbose:
print(f"\n📊 Step 4 - 重排序:")
print(
f" 重新排序后的最高分数 : {reranked_results[0].score:.3f}" if reranked_results else " 无结果")
else:
reranked_results = fused_results
if verbose:
print(f"\n📊 Step 4 - 重排序: 禁用了")
# ============================================================
# Step 5: 上下文压缩
# ============================================================
if enable_compression:
compressed_docs = self.compressor.compress(query, reranked_results)
if verbose:
print(f"\n📦 Step 5 - 上下文压缩:")
print(f" 原始文档长度: {sum(len(d.content) for d in reranked_results[:2])} 字符")
print(f" 压缩后文档长度: {sum(len(d.content) for d in compressed_docs[:2])} 字符")
else:
compressed_docs = reranked_results
if verbose:
print(f"\n📦 Step 5 - 上下文压缩：禁用了")
# ============================================================
# Step 6: 生成答案
# ============================================================
if verbose:
print(f"\n🤖 Step 6 - 生成答案:")
answer = self.generator.generate(query, compressed_docs)
if verbose:
print(f" 生成的答案长度: {len(answer)} 字符")
# ============================================================
# 返回结果
# ============================================================
result = {
"query": query,
"rewritten_queries": queries if enable_rewrite else [query],
"retrieved_count": total_retrieved,
"fused_count": len(fused_results) if enable_fusion else total_retrieved,
"final_documents": compressed_docs[:5], # 只返回前5个
"answer": answer
}
if verbose:
print(f"\n{'=' * 60}")
print(f"✅ Final Answer:")
print(f"{'=' * 60}")
print(answer)
print(f"{'=' * 60}\n")
return result
# =============================================================================
# 第十部分：演示和测试，完整的使用流程与测试
# =============================================================================
def demo():
"""
演示 Modular RAG 系统的完整使用流程
"""
print("=" * 70)
print("Modular RAG（模块化检索增强生成）演示程序")
print("=" * 70)
# 1. 准备文档数据
documents = [
"Python 是一种高级、解释型的编程语言。它由 Guido van Rossum 创建，于 1991 年首次发布。Python 强调代码可读性和简洁性。",
"RAG 检索增强生成是 Retrieval-Augmented Generation 的缩写。它将信息检索系统与大语言模型相结合，以生成更准确的回答。",
"机器学习 Machine Learning 是人工智能的一个子集。机器学习使系统能够从经验中学习和改进，而无需显式编程。机器学习的基本概念包括监督学习、无监督学习和强化学习。",
"Transformer 变换器架构由 Vaswani 等人在 2017 年发表的论文《Attention Is All You Need》中提出。Transformer 彻底改变了自然语言处理领域。",
"ChatGPT 对话生成预训练模型是由 OpenAI 开发的对话式人工智能模型。ChatGPT 基于 GPT（生成式预训练 Transformer）架构。",
"向量检索 Vector Search 是一种基于向量表示来查找相似项的技术。向量检索常用于推荐系统和信息检索中。",
"BM25 最佳匹配25是一种排序函数，搜索引擎用它来估算文档与给定查询的相关性。BM25 基于概率检索框架。",
"HyDE 假设文档嵌入是一种从查询生成假设文档并用它进行检索的技术。",
"重排序 Reranking 是使用更精确但计算成本更高的模型重新排列检索结果的过程，可以提升最终的排序质量。",
"上下文压缩 Context Compression 通过只提取最相关的句子来减少检索文档的长度，从而在大语言模型的上下文窗口中容纳更多信息。"
]
# 2. 创建 RAG 系统
print("\n📁 第一步：创建 Modular RAG 系统")
print("-" * 50)
rag = ModularRAG(config={
"fusion_method": "rrf", # 使用 RRF（倒排序） 融合
"max_sentences": 3, # 每个文档保留3个句子
"enable_rewrite": True, # 启用查询改写
"enable_rerank": True, # 启用重排序
"enable_fusion": True, # 启用融合
"enable_compression": True, # 启用压缩
"rewrite_method": "multi" # 多查询改写
})
# 3. 索引文档
print("\n📁 第二步：索引文档")
print("-" * 50)
rag.index_documents(documents)
# 4. 执行查询，输入关键词
test_queries = [
"什么是 RAG？",
"介绍一下 Python 编程语言",
"机器学习的基本概念是什么",
"什么是 BM25，它是如何工作的？"
]
results = {}
for query in test_queries:
result = rag.query(query, verbose=True)
results[query] = result
return results
# =======================================
# 主程序入口
# =======================================
if __name__ == "__main__":
# 运行演示
demo_results = demo()
print("\n" + "=" * 70)
print("🎉 Demo completed successfully!")
print("=" * 70)
运行效果：
好啦，更多内容请持续关注本站
更多推荐
8
0- 0
已为社区贡献1条内容
所有评论(0)