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

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

import sys
import os
import time
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from foundation.ai.rag.retrieval.retrieval import retrieval_manager
from foundation.observability.logger.loggering import server_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_entity"
    query = "起重机械 "
    
    # 测试多路召回
    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()