Genarrative/tools/test_ve_api.py

"""
Vector Engine API 连通性与生图耗时测试脚本。

用法:
  python tools/test_ve_api.py
  python tools/test_ve_api.py --prompt "你的自定义提示词"
  python tools/test_ve_api.py --size 1024x1024 --samples 3

前置条件:
  环境变量 VECTOR_ENGINE_API_KEY 已设置，
  或从 ../.env.secrets.local 自动读取。
"""

import os
import sys
import time
import json
import argparse
import requests
from concurrent.futures import ThreadPoolExecutor, as_completed


API_URL = "https://api.vectorengine.cn/v1/images/generations"
DEFAULT_PROMPT = "生成一张白色背景上的一只飞踢橘猫，绘本风格，不要文字水印"
DEFAULT_SIZE = "1024x1536"
DEFAULT_NEGATIVE = "文字、Logo、水印、按钮、UI、网格线、边框、编号、标签、纯色背景、白底、孤立主体"


def load_env_from_file(filepath):
    """从 .env 文件中加载环境变量（简单实现）"""
    if not os.path.exists(filepath):
        return
    with open(filepath, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            if "=" in line:
                key, _, value = line.partition("=")
                key = key.strip()
                value = value.strip().strip('"').strip("'")
                if key and value and key not in os.environ:
                    os.environ[key] = value


def single_request(api_key, base_url, prompt, negative, size, quality, index):
    """单次生图请求，返回 (耗时秒, task_id, 图片字节数)"""
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }
    body = {
        "model": "gpt-image-2",
        "prompt": prompt,
        "n": 1,
        "size": size,
    }
    if negative:
        body["negative_prompt"] = negative
    if quality:
        body["quality"] = quality

    start = time.time()
    try:
        resp = requests.post(
            base_url.rstrip("/") + "/v1/images/generations",
            headers=headers,
            json=body,
            timeout=600,
        )
        elapsed = time.time() - start

        if resp.status_code != 200:
            print(f"  [#{index}] HTTP {resp.status_code}: {resp.text[:300]}")
            return elapsed, None, 0

        data = resp.json()
        task_id = data.get("task_id", "")
        images = data.get("data", [])
        b64_len = len(images[0].get("b64_json", "")) if images else 0
        url = images[0].get("url", "") if images else ""

        print(f"  [#{index}] {elapsed:.1f}s  task_id={task_id}  b64={b64_len}chars  url={'present' if url else 'none'}")
        return elapsed, task_id, b64_len
    except requests.Timeout:
        elapsed = time.time() - start
        print(f"  [#{index}] TIMEOUT after {elapsed:.0f}s")
        return elapsed, None, 0
    except Exception as e:
        elapsed = time.time() - start
        print(f"  [#{index}] ERROR: {e}")
        return elapsed, None, 0


def main():
    parser = argparse.ArgumentParser(description="Vector Engine API 测试")
    parser.add_argument("--prompt", default=DEFAULT_PROMPT, help="生图提示词")
    parser.add_argument("--negative", default=DEFAULT_NEGATIVE, help="负面提示词")
    parser.add_argument("--size", default=DEFAULT_SIZE, help="图片尺寸 (1024x1024 / 1024x1536 / 1536x1024)")
    parser.add_argument("--samples", type=int, default=1, help="请求次数")
    parser.add_argument("--parallel", type=int, default=1, help="并行请求数 (默认1=串行)")
    parser.add_argument("--quality", default="", help="生图质量 (low/medium/high)")
    parser.add_argument("--base-url", default=API_URL, help="API 地址")
    args = parser.parse_args()

    # 自动加载 secrets
    script_dir = os.path.dirname(os.path.abspath(__file__))
    repo_root = os.path.dirname(script_dir)
    for fname in [".env.secrets.local", ".env.local", ".env"]:
        load_env_from_file(os.path.join(repo_root, fname))

    api_key = os.environ.get("VECTOR_ENGINE_API_KEY", "")
    if not api_key:
        print("错误: 未设置 VECTOR_ENGINE_API_KEY")
        print("请设置环境变量或将密钥写入 .env.secrets.local")
        sys.exit(1)

    base_url = os.environ.get("VECTOR_ENGINE_BASE_URL", args.base_url)
    print(f"API:   {base_url}")
    print(f"Size:  {args.size}")
    print(f"Samples: {args.samples}")
    print(f"Parallel: {args.parallel}")
    if args.quality:
        print(f"Quality: {args.quality}")
    print(f"Prompt ({len(args.prompt)} chars):")
    print(f"  {args.prompt[:120]}...")
    print(f"Negative ({len(args.negative)} chars):")
    print(f"  {args.negative[:120]}...")
    print()

    parallel = args.parallel
    total_start = time.time()

    if parallel <= 1:
        times = []
        for i in range(1, args.samples + 1):
            elapsed, task_id, b64_len = single_request(
                api_key, base_url, args.prompt, args.negative, args.size, args.quality, i
            )
            if b64_len > 0:
                times.append(elapsed)
    else:
        times = []
        with ThreadPoolExecutor(max_workers=parallel) as pool:
            futures = {
                pool.submit(
                    single_request,
                    api_key, base_url, args.prompt, args.negative, args.size, args.quality, idx
                ): idx
                for idx in range(1, args.samples + 1)
            }
            for future in as_completed(futures):
                elapsed, task_id, b64_len = future.result()
                if b64_len > 0:
                    times.append(elapsed)

    total_elapsed = time.time() - total_start

    if times:
        avg = sum(times) / len(times)
        print(f"\n成功: {len(times)}/{args.samples}")
        print(f"总耗时:  {total_elapsed:.1f}s")
        print(f"平均:    {avg:.1f}s")
        print(f"最快:    {min(times):.1f}s")
        print(f"最慢:    {max(times):.1f}s")
    else:
        print(f"\n全部失败 ({args.samples} 次)" + f" | 总耗时: {total_elapsed:.1f}s")


if __name__ == "__main__":
    main()