test_rag.py 991 B

12345678910111213141516171819202122232425
  1. import os
  2. import sys
  3. current_script_path = os.path.abspath(__file__)
  4. script_dir = os.path.dirname(current_script_path)
  5. project_root = os.path.abspath(os.path.join(script_dir, "../../"))
  6. if project_root not in sys.path:
  7. from foundation.ai.rag.retrieval.retrieval import retrieval_manager
  8. query = "实体:通风孔;背景:在箱梁腹板上设置通风孔,用于结构内部空气流通,需保证与预应力钢筋保护层满足间距要求,并水平设置"
  9. collection = "first_bfp_collection" # 你的 Milvus 集合名
  10. # 二阶段:Milvus混合检索(向量+BM25) → BGE重排
  11. results = retrieval_manager.multi_stage_recall(
  12. collection_name=collection,
  13. query_text=query,
  14. hybrid_top_k=50, # 第一阶段取多少候选
  15. top_k=3, # 最终返回条数
  16. ranker_type="weighted" # 或 "rrf"
  17. )
  18. print(results)
  19. for item in results:
  20. print(item["rerank_score"], item["text_content"])
  21. # 元数据在 item["metadata"](来自混合检索阶段)