LQAgentPlatform/utils_test/RAG_Test/test_rag.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
测试多阶段召回功能
"""

import sys
import os
import time

from foundation.ai.rag.retrieval.retrieval import retrieval_manager
from foundation.observability.logger.loggering import review_logger as logger


def test_multi_stage_recall(collection_name,query):
    """
    测试多阶段召回
    """

    try:
        start_time = time.time()
        results = retrieval_manager.multi_stage_recall(
            collection_name=collection_name,
            query_text=query,
            hybrid_top_k=10,
            top_k=5,
        )
        logger.info(f"返回结果:{results}")
        end_time = time.time()
        elapsed_time = end_time - start_time

        print(f"[OK] 召回完成，耗时: {elapsed_time:.2f}秒")
        print(f"[OK] 返回结果数量: {len(results)}")


    except Exception as e:
        print(f"[ERROR] 多阶段召回测试失败: {str(e)}")

def test_hybrid_search_recall(collection_name,query):
    """
    测试混合召回
    """
    try:
        start_time = time.time()
        results = retrieval_manager.hybrid_search_recall(
            collection_name=collection_name,
            query_text=query,
            top_k=1,
            ranker_type="weighted",
            dense_weight=0.7,
            sparse_weight=0.3
        )
        logger.info(f"返回结果:{results}")
        end_time = time.time()
        elapsed_time = end_time - start_time
        print(f"[OK] 召回完成，耗时: {elapsed_time:.2f}秒")
        print(f"[OK] 召回结果数量: {len(results)}")
        return results
    
    except Exception as e:
        print(f"[ERROR] 混合召回测试失败: {str(e)}")




def main():
    """
    主测试函数
    """
    collection_name = "first_bfp_collection_test"
    query = "起重小车轨道，起重量小于 320t的分段拼接桁架梁每段梁上小车轨道不允许有接缝（允许焊为一体），拼接 处高低差≤2mm、间隙≤4mm、侧向错位≤2mm，非焊接连接轨道端部加挡铁,其他梁轨道接头高低差≤1mm、间隙≤2mm、侧向错位≤1mm，正轨箱形梁及半偏轨箱形梁轨道接缝应放 在筋板上允差≤15mm，两端最短轨道长度≥1.5m且端部加挡"
    
    # 测试多路召回
    logger.info("开始测试多路召回...")
    test_multi_stage_recall(collection_name,query=query)

    # # 测试混合召回
    # logger.info("开始测试混合召回...")
    # test_hybrid_search_recall(collection_name="first_bfp_collection_entity",query=query)





if __name__ == "__main__":
    main()