SpinachLove/code-rag-assistant

1from rag.vector_db import similarity_search
2from rag.rerank import rerank_documents
3
4# 设计10个测试问题，覆盖你的知识库内容，分为基础题、细节题、场景题
5test_questions = [
6    "Python的核心数据类型有哪些？",
7    "Python如何处理异常？try-except的用法是什么？",
8    "什么是Python的装饰器？怎么使用？",
9    "Python的深拷贝和浅拷贝有什么区别？",
10    "列表推导式的用法和优势是什么？",
11    "Python中的self参数有什么作用？",
12    "什么是递归函数？递归的优缺点是什么？",
13    "Python的面向对象三大特性是什么？",
14    "如何解决Python的索引越界问题？",
15    "Python中的生成器和迭代器有什么区别？"
16]
17
18# 人工标注的正确答案关键词，用来判断检索结果是否相关
19# 格式：问题索引 -> 必须包含的关键词列表
20standard_keywords = {
21    0: ["列表", "元组", "可变", "不可变", "修改"],
22    1: ["try", "except", "异常", "捕获", "报错"],
23    2: ["装饰器", "函数", "包装", "@", "语法糖"],
24    3: ["深拷贝", "浅拷贝", "引用", "地址", "copy"],
25    4: ["列表推导式", "循环", "简洁", "语法", "遍历"],
26    5: ["self", "实例", "类", "方法", "对象"],
27    6: ["递归", "函数", "自身", "终止条件", "栈"],
28    7: ["封装", "继承", "多态", "面向对象", "类"],
29    8: ["索引", "越界", "长度", "下标", "IndexError"],
30    9: ["生成器", "迭代器", "yield", "next", "可迭代"]
31}
32
33def calculate_hit_rate(docs, keywords):
34    """计算检索结果的命中率，只要有一个文档包含所有关键词，就算命中"""
35    for doc in docs:
36        content = doc["content"].lower()
37        all_keywords_hit = all(keyword.lower() in content for keyword in keywords)
38        if all_keywords_hit:
39            return True
40    return False
41
42def calculate_top3_accuracy(docs, keywords):
43    """计算top3结果的准确率，前3条里有相关内容就算准确"""
44    return calculate_hit_rate(docs[:3], keywords)
45
46if __name__ == "__main__":
47    print("===== RAG优化效果对比测试 =====")
48    print(f"测试问题数量：{len(test_questions)}个")
49    print("-" * 50)
50
51    # 统计优化前（仅向量检索）的准确率
52    before_optimization_correct = 0
53    # 统计优化后（向量检索+重排序）的准确率
54    after_optimization_correct = 0
55
56    # 逐个问题测试
57    for idx, question in enumerate(test_questions):
58        print(f"\n测试问题{idx+1}：{question}")
59        keywords = standard_keywords.get(idx, [])
60        if not keywords:
61            print("跳过：无标注关键词")
62            continue
63
64        # 优化前：仅向量检索top3
65        retrieve_docs = similarity_search(question, top_k=3)
66        before_correct = calculate_top3_accuracy(retrieve_docs, keywords)
67        if before_correct:
68            before_optimization_correct += 1
69        print(f"优化前（仅向量检索top3）：{'准确' if before_correct else '不准确'}")
70
71        # 优化后：向量检索top10 + 重排序top3
72        full_retrieve_docs = similarity_search(question, top_k=10)
73        reranked_docs = rerank_documents(question, full_retrieve_docs, top_n=3)
74        after_correct = calculate_top3_accuracy(reranked_docs, keywords)
75        if after_correct:
76            after_optimization_correct += 1
77        print(f"优化后（检索+重排序top3）：{'准确' if after_correct else '不准确'}")
78
79    # 计算最终准确率
80    total_test = len(standard_keywords)           #总测试关键字个数
81    before_accuracy = (before_optimization_correct / total_test) * 100
82    after_accuracy = (after_optimization_correct / total_test) * 100
83    promotion = after_accuracy - before_accuracy
84
85    print("\n" + "=" * 50)
86    print("===== 最终测试结果 =====")
87    print(f"总测试用例数：{total_test}个")
88    print(f"优化前准确率：{before_accuracy:.2f}%")
89    print(f"优化后准确率：{after_accuracy:.2f}%")
90    print(f"准确率提升：{promotion:.2f}个百分点")
91    print("=" * 50)
92
93    # 生成测试报告内容
94    report_content = f"""
95    # RAG优化效果测试报告
96    ## 测试背景
97    为验证Re-rank重排序技术对RAG检索精度的优化效果，设计了{total_test}个编程领域测试问题，对比优化前后的top3检索准确率。
98
99    ## 测试方案
100    - 优化前：仅使用向量相似度检索，返回top3结果
101    - 优化后：先通过向量检索完成top10粗召回，再通过Cohere Re-rank模型完成top3精排
102    - 评估标准：top3检索结果中包含问题的核心知识点，即为准确
103
104    ## 测试结果
105    | 指标 | 优化前 | 优化后 | 提升幅度 |
106    |------|--------|--------|----------|
107    | 准确率 | {before_accuracy:.2f}% | {after_accuracy:.2f}% | {promotion:.2f}% |
108
109    ## 结论
110    引入Re-rank重排序技术后，RAG系统的检索准确率提升了{promotion:.2f}个百分点，有效解决了向量检索只关注语义相似度、忽略关键词匹配的问题，大幅提升了问答的准确性，同时降低了大模型的幻觉风险。
111    """
112
113    # 把测试报告写入docs文件夹
114    with open("../docs/RAG优化效果测试报告.md", "w", encoding="utf-8") as f:
115        f.write(report_content)
116    print("\n测试报告已生成：docs/RAG优化效果测试报告.md")