four-level-schedule/test_multimodal_search.py

#!/usr/bin/env python3
"""
测试多模态检索 API
"""
import requests
import json

# 配置
EMBEDDING_API_URL = "http://localhost:8084"


def test_text_only_search():
    """测试纯文本检索（兼容现有能力）"""
    print("=" * 60)
    print("测试 1: 纯文本检索 (TEXT_ONLY)")
    print("=" * 60)

    url = f"{EMBEDDING_API_URL}/api/v2/multimodal/search"

    request_body = {
        "searchType": "TEXT_ONLY",
        "textQuery": "数学",
        "topK": 5
    }

    try:
        response = requests.post(
            url,
            json=request_body,
            headers={"Content-Type": "application/json"},
            timeout=30,
        )

        print(f"Status: {response.status_code}")
        result = response.json()
        print(f"Response:")
        print(json.dumps(result, ensure_ascii=False, indent=2))

        if result.get("success"):
            print(f"\n✅ 成功！找到 {len(result.get('allResults', []))} 个结果")
            for i, item in enumerate(result.get('allResults', []), 1):
                print(f"\n  结果 {i}:")
                print(f"    文档ID: {item.get('docId')}")
                print(f"    相似度: {item.get('score'):.4f}")
                content = item.get('content', '')
                if len(content) > 80:
                    content = content[:80] + "..."
                print(f"    内容: {content}")
        else:
            print(f"\n❌ 失败: {result.get('errorMessage')}")

    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()

    print()


def test_text_to_multimodal_search():
    """测试文本查多模态"""
    print("=" * 60)
    print("测试 2: 文本查多模态 (TEXT_TO_MULTIMODAL)")
    print("=" * 60)

    url = f"{EMBEDDING_API_URL}/api/v2/multimodal/search"

    request_body = {
        "searchType": "TEXT_TO_MULTIMODAL",
        "textQuery": "口算题",
        "topK": 5
    }

    try:
        response = requests.post(
            url,
            json=request_body,
            headers={"Content-Type": "application/json"},
            timeout=30,
        )

        print(f"Status: {response.status_code}")
        result = response.json()
        print(f"Response:")
        print(json.dumps(result, ensure_ascii=False, indent=2))

    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()

    print()


def test_image_to_multimodal_search():
    """测试图片查多模态（待实现）"""
    print("=" * 60)
    print("测试 3: 图片查多模态 (IMAGE_TO_MULTIMODAL) - 待实现")
    print("=" * 60)

    url = f"{EMBEDDING_API_URL}/api/v2/multimodal/search"

    request_body = {
        "searchType": "IMAGE_TO_MULTIMODAL",
        "imageData": "base64_encoded_image_here",
        "imageDataType": "BASE64",
        "topK": 5
    }

    try:
        response = requests.post(
            url,
            json=request_body,
            headers={"Content-Type": "application/json"},
            timeout=30,
        )

        print(f"Status: {response.status_code}")
        result = response.json()
        print(f"Response:")
        print(json.dumps(result, ensure_ascii=False, indent=2))

    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()

    print()


def main():
    print("\n")
    print("╔" + "=" * 58 + "╗")
    print("║" + " " * 15 + "多模态检索 API 测试" + " " * 18 + "║")
    print("╚" + "=" * 58 + "╝")
    print()

    # 检查服务状态
    print("检查服务状态...")
    try:
        health_response = requests.get(f"{EMBEDDING_API_URL}/actuator/health", timeout=5)
        if health_response.status_code == 200:
            health_data = health_response.json()
            if health_data.get("status") == "UP":
                print("✅ schedule-embedding-api 正常\n")
            else:
                print(f"⚠️  schedule-embedding-api 状态: {health_data.get('status')}\n")
        else:
            print(f"❌ schedule-embedding-api 异常: {health_response.status_code}\n")
            return
    except Exception as e:
        print(f"❌ schedule-embedding-api 不可达: {e}\n")
        return

    # 运行测试
    test_text_only_search()
    test_text_to_multimodal_search()
    test_image_to_multimodal_search()

    print("=" * 60)
    print("测试完成")
    print("=" * 60)


if __name__ == "__main__":
    main()