import sys import os sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src")) from base_rag import BaseRAG class SimpleRAG(BaseRAG): def ingest(self, documents): for doc in documents: self.vector_store.add_texts([doc]) def query(self, question, k=3): # 使用带重排功能的搜索 print(f"查询参数 k={k}") docs = self.similarity_search_with_rerank(question, k=k) print(f"重排后返回文档数量: {len(docs)}") return docs def query_without_rerank(self, question, k=3): # 不使用重排的普通搜索 docs = self.similarity_search(question, k=k) return docs def main(): # 重排配置 - 使用BGE重排模型 rerank_config = { "enabled": True, "type": "local", "model": "BAAI/bge-reranker-base", # 注意:这里不设置top_k,让它在查询时动态决定 } rag = SimpleRAG(rerank_config=rerank_config) print("RAG系统(含重排功能)初始化完成!") # 添加更多测试文档 # documents = [ # "苹果是一种红色或绿色的水果,味道甜美,营养丰富,含有丰富的维生素C。", # "苹果公司是一家总部位于美国的科技公司,以生产iPhone、iPad、Mac等产品而闻名。", # "Python是一种高级编程语言,简单易学,功能强大,广泛用于数据科学和机器学习。", # "苹果树是一种果树,春天开花,秋天结果,需要充足的阳光和水分。", # "苹果派是一种传统的美式甜点,由苹果馅和酥脆的派皮制成。", # "苹果醋是由苹果发酵制成的,具有一定的保健功效,可以帮助消化。", # "iPhone是苹果公司生产的智能手机,具有先进的技术和优秀的用户体验。", # "机器学习是人工智能的一个分支,Python是机器学习领域最流行的编程语言之一。", # "苹果公司主要产品是iPhone、iPad、Mac等", # "使用多线程可以提高程序在多核处理器上的运行效率。", # "数据库设计应该遵循第三范式,以减少数据冗余。", # "使用更高效的数据结构(如哈希表)可以提升算法效率。", # "程序界面应该简洁明了,方便用户操作。", # "尽量避免使用全局变量,减少潜在的并发问题。", # "在Python中使用NumPy等库可以加快数值计算速度。", # "使用SSD替代传统HDD可提升文件读写速度。", # "将代码中的重复逻辑封装为函数,有助于维护但对性能影响小。", # "内存泄漏会导致程序运行越来越慢,应及时释放资源。", # "使用代码缓存(如Memoization)可以避免重复计算。", # ] # print("正在添加文档...") # rag.ingest(documents) # print(f"文档添加完成! 共添加了 {len(documents)} 个文档") # 测试查询并比较重排效果 test_query = "Python中的list操作较慢,尽量使用数组库" print(f"\n测试查询: '{test_query}'") print("\n=== 不使用重排的结果 ===") result_no_rerank = rag.query_without_rerank(test_query, k=3) for i, doc in enumerate(result_no_rerank, 1): print(f"{i}. {doc.page_content}") print("\n=== 使用重排的结果 ===") result_with_rerank = rag.query(test_query, k=3) for i, doc in enumerate(result_with_rerank, 1): print(f"{i}. {doc.page_content}") if __name__ == "__main__": main()