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

"""
简化的重排序测试脚本
"""

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

from foundation.ai.models.rerank_model import execute_rerank

def test_rerank_function():
    """
    测试重排序功能
    """
    print("开始测试重排序功能")
    print("=" * 50)

    # 测试用例1：乔布斯查询
    query1 = "乔布斯是谁？"
    candidates1 = [
        "大模型是一类具有大量参数的人工智能模型。",
        "苹果是一家科技公司",
        "大模型用于深度学习任务"
    ]

    print("测试用例1:")
    print(f"查询: {query1}")
    print("候选文档:")
    for i, doc in enumerate(candidates1, 1):
        print(f"  {i}. {doc}")

    try:
        results1 = execute_rerank(query1, candidates1, top_k=3)

        print("\n重排序结果:")
        for i, result in enumerate(results1, 1):
            text = result.get("text", "")
            score = result.get("score", "0.0")
            print(f"  {i}. [{score}] {text}")

        print(f"\n结果数量: {len(results1)}")
        if results1:
            print("测试用例1: 成功")
        else:
            print("测试用例1: 失败 - 没有返回结果")

    except Exception as e:
        print(f"测试用例1: 失败 - 异常: {str(e)}")

    print("\n" + "-" * 50)

    # 测试用例2：人工智能查询
    query2 = "什么是人工智能？"
    candidates2 = [
        "人工智能是计算机科学的一个分支。",
        "机器学习是人工智能的核心技术。",
        "深度学习使用神经网络进行学习。",
        "苹果公司生产iPhone手机。",
        "Transformer是大模型的基础架构。"
    ]

    print("测试用例2:")
    print(f"查询: {query2}")
    print("候选文档:")
    for i, doc in enumerate(candidates2, 1):
        print(f"  {i}. {doc}")

    try:
        results2 = execute_rerank(query2, candidates2, top_k=5)

        print("\n重排序结果:")
        for i, result in enumerate(results2, 1):
            text = result.get("text", "")
            score = result.get("score", "0.0")
            print(f"  {i}. [{score}] {text}")

        print(f"\n结果数量: {len(results2)}")
        if results2:
            print("测试用例2: 成功")
        else:
            print("测试用例2: 失败 - 没有返回结果")

    except Exception as e:
        print(f"测试用例2: 失败 - 异常: {str(e)}")

    print("\n" + "=" * 50)
    print("重排序功能测试完成")

if __name__ == "__main__":
    test_rerank_function()