LQAgentPlatform/utils_test/RAG_Test/test_entity_bfp_recall.py

import json
import asyncio
from foundation.ai.rag.retrieval.retrieval import retrieval_manager
from foundation.observability.monitoring.time_statistics import track_execution_time


entity = "架桥机"
search_keywords = ["提梁机", "架桥设备", "造桥机"]
background = "JQ220t-40m架桥机安装及拆除"

@track_execution_time
def main():

    entity_list = asyncio.run(retrieval_manager.entity_recall(entity,search_keywords,top_k=5))
    bfp_result = asyncio.run(retrieval_manager.async_bfp_recall(entity_list,background,top_k=10))
    with open("temp\entity_bfp_recall\entity_bfp_recall.json", "w", encoding="utf-8") as f:
        json.dump(bfp_result, f, ensure_ascii=False, indent=4)

main()