LQAgentPlatform/test/test_rerank_api_curl.py

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

"""
使用curl测试重排序API
"""

import subprocess
import json
import sys
import os

def test_rerank_api_with_curl():
    """
    使用curl命令测试重排序API
    """
    print("=== 使用curl测试重排序API ===")

    # 构建curl命令
    curl_command = [
        "curl",
        "--location",
        "http://192.168.91.253:9005/v1/rerank",
        "--header", "Content-Type: application/json",
        "--data", json.dumps({
            "model": "bge-reranker-v2-m3",
            "query": "乔布斯是谁？",
            "candidates": [
                "大模型是一类具有大量参数的人工智能模型。",
                "苹果是一家科技公司",
                "大模型用于深度学习任务"
            ]
        })
    ]

    try:
        print("执行命令:")
        print(" ".join(curl_command))
        print()

        # 执行curl命令
        result = subprocess.run(
            curl_command,
            capture_output=True,
            text=True,
            timeout=30
        )

        print(f"返回状态码: {result.returncode}")
        print("=" * 50)
        print("标准输出:")
        print(result.stdout)

        if result.stderr:
            print("=" * 50)
            print("标准错误:")
            print(result.stderr)

        # 尝试解析JSON响应
        if result.stdout:
            try:
                response_data = json.loads(result.stdout)
                print("=" * 50)
                print("解析后的JSON响应:")
                print(json.dumps(response_data, ensure_ascii=False, indent=2))

                # 检查响应格式
                if "results" in response_data:
                    print("\n✓ API响应格式正确，包含results字段")
                    results = response_data["results"]
                    print(f"✓ 返回了 {len(results)} 个重排序结果")

                    for i, item in enumerate(results, 1):
                        text = item.get("text", "")
                        score = item.get("score", "")
                        print(f"  {i}. 分数: {score} | 内容: {text}")
                else:
                    print("\n✗ API响应格式异常，缺少results字段")

            except json.JSONDecodeError as e:
                print(f"\n✗ JSON解析失败: {str(e)}")
        else:
            print("\n✗ 没有收到任何响应")

        return result.returncode == 0

    except subprocess.TimeoutExpired:
        print("✗ 请求超时")
        return False
    except Exception as e:
        print(f"✗ 执行curl命令时发生异常: {str(e)}")
        return False

def test_different_queries():
    """
    测试不同的查询请求
    """
    print("\n=== 测试不同的查询请求 ===")

    test_cases = [
        {
            "query": "什么是人工智能？",
            "candidates": [
                "人工智能是计算机科学的一个分支。",
                "机器学习是人工智能的核心技术。",
                "深度学习使用神经网络进行学习。"
            ]
        },
        {
            "query": "大模型有什么特点？",
            "candidates": [
                "大模型具有数百万到数十亿的参数。",
                "苹果公司生产iPhone手机。",
                "Transformer是大模型的基础架构。"
            ]
        },
        {
            "query": "机器学习和深度学习的区别",
            "candidates": [
                "深度学习是机器学习的一个子集。",
                "机器学习需要人工特征工程。",
                "深度学习可以自动学习特征。"
            ]
        }
    ]

    for i, test_case in enumerate(test_cases, 1):
        print(f"\n--- 测试用例 {i}: {test_case['query']} ---")

        curl_command = [
            "curl",
            "--location",
            "http://192.168.91.253:9005/v1/rerank",
            "--header", "Content-Type: application/json",
            "--data", json.dumps({
                "model": "bge-reranker-v2-m3",
                "query": test_case["query"],
                "candidates": test_case["candidates"]
            })
        ]

        try:
            result = subprocess.run(
                curl_command,
                capture_output=True,
                text=True,
                timeout=30
            )

            if result.returncode == 0 and result.stdout:
                try:
                    response_data = json.loads(result.stdout)
                    if "results" in response_data:
                        results = response_data["results"]
                        print(f"✓ 成功返回 {len(results)} 个结果")

                        # 显示前3个结果
                        for j, item in enumerate(results[:3], 1):
                            text = item.get("text", "")
                            score = item.get("score", "")
                            print(f"  {j}. [{score}] {text[:50]}...")
                    else:
                        print("✗ 响应格式异常")
                except json.JSONDecodeError:
                    print("✗ JSON解析失败")
            else:
                print(f"✗ 请求失败: 状态码 {result.returncode}")

        except Exception as e:
            print(f"✗ 测试失败: {str(e)}")

def main():
    """
    主函数
    """
    print("开始使用curl测试重排序API")
    print("=" * 60)

    # 基本测试
    success = test_rerank_api_with_curl()

    if success:
        print("\n🎉 基本API测试成功！")

        # 测试不同查询
        test_different_queries()

        print("\n" + "=" * 60)
        print("所有测试完成！")

    else:
        print("\n❌ 基本API测试失败！")
        print("请检查:")
        print("1. 重排序服务是否在 192.168.91.253:9005 运行")
        print("2. 服务是否支持 /v1/rerank 端点")
        print("3. 网络连接是否正常")

        return 1

    return 0

if __name__ == "__main__":
    sys.exit(main())