DataOps-platform/scripts/test_cohere_api_key.py

"""
测试 Cohere API Key 是否有效

用于验证 API Key 是否可以正常使用
"""

import sys
import os
from typing import Optional

import requests
from loguru import logger

# 配置日志
logger.remove()
logger.add(
    sys.stdout,
    level="INFO",
    format="<green>{time:YYYY-MM-DD HH:mm:ss}</green> | <level>{level: <8}</level> | <level>{message}</level>",
)


def test_cohere_api_key(api_key: str) -> dict:
    """
    测试 Cohere API Key 是否有效

    Args:
        api_key: Cohere API Key

    Returns:
        测试结果
    """
    # Cohere API 端点
    base_url = "https://api.cohere.ai/v1"

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    # 测试 1: 检查 API Key 基本信息
    logger.info("测试 1: 检查 API Key 基本信息...")
    try:
        # 使用 models 端点测试（轻量级）
        response = requests.get(
            f"{base_url}/models",
            headers=headers,
            timeout=10,
        )

        logger.debug(f"响应状态码: {response.status_code}")
        logger.debug(f"响应内容: {response.text[:200]}")

        if response.status_code == 200:
            logger.success("✅ API Key 验证成功！")
            models = response.json()
            return {
                "success": True,
                "message": "API Key 有效",
                "status_code": response.status_code,
                "models_count": len(models.get("models", [])),
            }
        elif response.status_code == 401:
            logger.error("❌ API Key 认证失败 - 无效的 API Key")
            return {
                "success": False,
                "message": "API Key 认证失败",
                "error": "Unauthorized",
                "status_code": 401,
            }
        elif response.status_code == 403:
            logger.error("❌ API Key 权限不足 - 可能是账户权限问题")
            return {
                "success": False,
                "message": "API Key 权限不足",
                "error": "Forbidden",
                "status_code": 403,
            }
        else:
            logger.warning(f"⚠️  意外的响应状态码: {response.status_code}")
            return {
                "success": False,
                "message": f"意外的响应: {response.status_code}",
                "error": response.text[:200],
                "status_code": response.status_code,
            }

    except requests.exceptions.RequestException as e:
        logger.error(f"❌ 请求异常: {str(e)}")
        return {
            "success": False,
            "message": f"请求失败: {str(e)}",
            "error": str(e),
        }

    # 测试 2: 尝试使用 rerank 端点（如果测试1失败）
    logger.info("测试 2: 尝试使用 rerank 端点...")
    try:
        rerank_data = {
            "model": "rerank-multilingual-v3.0",
            "query": "test query",
            "documents": ["test document 1", "test document 2"],
            "top_n": 2,
        }

        response = requests.post(
            f"{base_url}/rerank",
            headers=headers,
            json=rerank_data,
            timeout=10,
        )

        logger.debug(f"Rerank 响应状态码: {response.status_code}")
        logger.debug(f"Rerank 响应内容: {response.text[:200]}")

        if response.status_code == 200:
            logger.success("✅ Rerank API 测试成功！")
            return {
                "success": True,
                "message": "Rerank API 可用",
                "status_code": response.status_code,
            }
        else:
            logger.warning(f"⚠️  Rerank API 测试失败: {response.status_code}")
            return {
                "success": False,
                "message": f"Rerank API 测试失败: {response.status_code}",
                "error": response.text[:200],
                "status_code": response.status_code,
            }

    except requests.exceptions.RequestException as e:
        logger.error(f"❌ Rerank 请求异常: {str(e)}")
        return {
            "success": False,
            "message": f"Rerank 请求失败: {str(e)}",
            "error": str(e),
        }


def main():
    """主函数"""
    # 从命令行参数或环境变量获取 API Key
    api_key = sys.argv[1] if len(sys.argv) > 1 else os.environ.get("COHERE_API_KEY")

    if not api_key:
        logger.error("❌ 请提供 Cohere API Key")
        logger.info("使用方法: python test_cohere_api_key.py <API_KEY>")
        logger.info("或设置环境变量: COHERE_API_KEY=<API_KEY>")
        return 1

    # 隐藏部分 API Key（安全显示）
    masked_key = api_key[:8] + "..." + api_key[-4:] if len(api_key) > 12 else "***"
    logger.info(f"🔍 测试 Cohere API Key: {masked_key}")

    result = test_cohere_api_key(api_key)

    print("\n" + "=" * 60)
    print("测试结果:")
    print("=" * 60)
    print(f"成功: {result.get('success', False)}")
    print(f"消息: {result.get('message', 'N/A')}")
    if result.get("error"):
        print(f"错误: {result.get('error')}")
    if result.get("status_code"):
        print(f"状态码: {result.get('status_code')}")

    if not result.get("success"):
        print("\n" + "=" * 60)
        print("可能的原因:")
        print("=" * 60)
        print("1. API Key 无效或已过期")
        print("2. API Key 格式不正确（应该没有空格或换行）")
        print("3. Cohere 账户权限不足")
        print("4. API Key 未激活或账户未验证")
        print("5. 网络连接问题")
        print("\n建议:")
        print("- 检查 Cohere Dashboard: https://dashboard.cohere.com/")
        print("- 确认 API Key 是否正确复制（无多余空格）")
        print("- 检查账户状态和权限")

    return 0 if result.get("success") else 1


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