test_dedup_api.py 10.6 KB
"""歌词去重 & 曲去重 API 测试脚本

用法:
    # ---- 歌词去重 ----
    # 上传指定歌词文件并调用去重 API
    python test_api/test_dedup_api.py --file data/library/None_WHHY134166.lrc

    # 指定标题和歌手
    python test_api/test_dedup_api.py --file data/library/None_WHHY134166.lrc --title "夜曲" --artist "周杰伦"

    # 仅上传不调用 API
    python test_api/test_dedup_api.py --file data/library/None_WHHY134166.lrc --upload-only

    # 仅调用 API(使用已有 URL)
    python test_api/test_dedup_api.py --url "https://hikoon-ai-test.oss-cn-hangzhou.aliyuncs.com/temp_ai/20250603/xxx.lrc"

    # 指定 API 地址
    python test_api/test_dedup_api.py --file data/library/None_WHHY134166.lrc --api-url "http://localhost:8000"

    # ---- 曲去重:入库 ----
    python test_api/test_dedup_api.py rhythm-ingest --song-id 12345 --file data/audio/song.mp3
    python test_api/test_dedup_api.py rhythm-ingest --song-id 12345 --url "https://example.com/song.mp3"

    # ---- 曲去重:查询 ----
    python test_api/test_dedup_api.py rhythm-check --file data/audio/query.mp3
    python test_api/test_dedup_api.py rhythm-check --url "https://example.com/query.mp3" --top-k 50
"""
import argparse
import json
import os
import sys

# 确保项目根目录在 path 中
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

import urllib.request
import urllib.error

from test_api.config import OSS_ACCESS_KEY_ID, OSS_ACCESS_KEY_SECRET, OSS_ENDPOINT, OSS_BUCKET_NAME
from test_api.oss_uploader import OSSUploader


def upload_lyric_file(file_path: str) -> str:
    """上传歌词文件到 OSS,返回公开 URL"""
    uploader = OSSUploader()
    success, result = uploader.upload_file(file_path)
    if not success:
        print(f"上传失败: {result}")
        sys.exit(1)
    return result


def upload_audio_file(file_path: str) -> str:
    """上传音频文件到 OSS,返回公开 URL"""
    uploader = OSSUploader()
    success, result = uploader.upload_file(file_path)
    if not success:
        print(f"上传失败: {result}")
        sys.exit(1)
    return result


def call_dedup_api(url: str, title: str | None, artist: str | None, api_base: str) -> dict:
    """调用去重 API"""
    payload = json.dumps({
        "url": url,
        "title": title,
        "artist": artist,
    }).encode("utf-8")

    req = urllib.request.Request(
        f"{api_base.rstrip('/')}/api/v1/check",
        data=payload,
        headers={"Content-Type": "application/json"},
        method="POST",
    )

    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            body = json.loads(resp.read().decode("utf-8"))
            return body
    except urllib.error.HTTPError as exc:
        error_body = exc.read().decode("utf-8", errors="replace")
        print(f"API 请求失败 (HTTP {exc.code}): {error_body}")
        sys.exit(1)
    except urllib.error.URLError as exc:
        print(f"API 请求失败: {exc.reason}")
        print("请确认 API 服务已启动: uvicorn lyric_dedup_server.app:app --host 0.0.0.0 --port 8000")
        sys.exit(1)


def main():
    parser = argparse.ArgumentParser(description="去重 API 测试")
    parser.add_argument("--api-url", default="http://localhost:8000", help="API 服务地址 (默认 http://localhost:8000)")

    subparsers = parser.add_subparsers(dest="command")

    # ---- 曲去重:入库 ----
    ingest_parser = subparsers.add_parser("rhythm-ingest", help="将音频特征写入曲库")
    ingest_parser.add_argument("--song-id", "-s", type=int, required=True, help="歌曲 ID")
    ingest_parser.add_argument("--file", "-f", help="本地音频文件路径")
    ingest_parser.add_argument("--url", "-u", help="已上传的音频 URL(跳过上传步骤)")

    # ---- 曲去重:查询 ----
    check_parser = subparsers.add_parser("rhythm-check", help="查询音频是否与曲库重复")
    check_parser.add_argument("--file", "-f", help="本地音频文件路径")
    check_parser.add_argument("--url", "-u", help="已上传的音频 URL(跳过上传步骤)")
    check_parser.add_argument("--top-k", type=int, default=100, help="召回候选数量(默认 100)")

    # ---- 歌词去重(默认,无子命令时走此路径) ----
    parser.add_argument("--file", "-f", help="本地歌词文件路径")
    parser.add_argument("--url", "-u", help="已上传的歌词 URL(跳过上传步骤)")
    parser.add_argument("--title", "-t", help="歌曲标题(可选)")
    parser.add_argument("--artist", "-a", help="歌手名(可选)")
    parser.add_argument("--upload-only", action="store_true", help="仅上传到 OSS,不调用 API")

    args = parser.parse_args()

    # ------------------------------------------------------------------ #
    # 曲去重:入库
    # ------------------------------------------------------------------ #
    if args.command == "rhythm-ingest":
        if not args.file and not args.url:
            ingest_parser.error("需要指定 --file 或 --url")

        if args.file:
            abs_path = os.path.join(PROJECT_ROOT, args.file) if not os.path.isabs(args.file) else args.file
            if not os.path.exists(abs_path):
                print(f"文件不存在: {abs_path}")
                sys.exit(1)
            print(f"正在上传音频: {abs_path}")
            audio_url = upload_audio_file(abs_path)
            print(f"上传成功: {audio_url}")
        else:
            audio_url = args.url
            print(f"使用已有 URL: {audio_url}")

        print(f"\n正在调用曲入库 API (song_id={args.song_id})...")
        result = _call_rhythm_ingest(audio_url, args.song_id, api_base=args.api_url)

        print(f"\n结果:")
        print(f"  song_id: {result.get('song_id')}")
        print(f"  success: {result.get('success')}")
        print(f"  message: {result.get('message', '')}")
        return

    # ------------------------------------------------------------------ #
    # 曲去重:查询
    # ------------------------------------------------------------------ #
    if args.command == "rhythm-check":
        if not args.file and not args.url:
            check_parser.error("需要指定 --file 或 --url")

        if args.file:
            abs_path = os.path.join(PROJECT_ROOT, args.file) if not os.path.isabs(args.file) else args.file
            if not os.path.exists(abs_path):
                print(f"文件不存在: {abs_path}")
                sys.exit(1)
            print(f"正在上传音频: {abs_path}")
            audio_url = upload_audio_file(abs_path)
            print(f"上传成功: {audio_url}")
        else:
            audio_url = args.url
            print(f"使用已有 URL: {audio_url}")

        print(f"\n正在调用曲去重 API...")
        result = _call_rhythm_check(audio_url, top_k=args.top_k, api_base=args.api_url)

        print(f"\n结果:")
        print(f"  duplicate:  {result.get('duplicate')}")
        candidates = result.get("candidates", [])
        if candidates:
            print(f"  candidates ({len(candidates)}):")
            for i, c in enumerate(candidates[:10]):
                print(f"    [{i+1}] song_id={c.get('song_id')}  similarity={c.get('similarity', 0):.4f}  source={c.get('source')}")
            if len(candidates) > 10:
                print(f"    ... 共 {len(candidates)} 条")
        else:
            print(f"  candidates: []")
        return

    # ------------------------------------------------------------------ #
    # 歌词去重(默认路径,无子命令)
    # ------------------------------------------------------------------ #
    if not args.file and not args.url:
        parser.error("需要指定 --file 或 --url,或使用子命令 rhythm-ingest / rhythm-check")

    # Step 1: 上传
    if args.file:
        abs_path = os.path.join(PROJECT_ROOT, args.file) if not os.path.isabs(args.file) else args.file
        if not os.path.exists(abs_path):
            print(f"文件不存在: {abs_path}")
            sys.exit(1)
        print(f"正在上传: {abs_path}")
        lyric_url = upload_lyric_file(abs_path)
        print(f"上传成功: {lyric_url}")
    else:
        lyric_url = args.url
        print(f"使用已有 URL: {lyric_url}")

    if args.upload_only:
        return

    # Step 2: 调用去重 API
    print(f"\n正在调用去重 API...")
    result = call_dedup_api(lyric_url, title=args.title, artist=args.artist, api_base=args.api_url)

    print(f"\n结果:")
    print(f"  duplicate:  {result.get('duplicate')}")
    print(f"  decision:   {result.get('decision', 'N/A')}")
    print(f"  confidence: {result.get('confidence', 'N/A')}")
    print(f"  reason:     {result.get('reason', 'N/A')}")
    print(f"  record_ids: {result.get('record_ids', [])}")


def _call_rhythm_ingest(url: str, song_id: int, api_base: str) -> dict:
    """调用曲入库 API"""
    payload = json.dumps({"song_id": song_id, "url": url}).encode("utf-8")
    req = urllib.request.Request(
        f"{api_base.rstrip('/')}/api/v1/rhythm/ingest",
        data=payload,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=120) as resp:
            return json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        error_body = exc.read().decode("utf-8", errors="replace")
        print(f"API 请求失败 (HTTP {exc.code}): {error_body}")
        sys.exit(1)
    except urllib.error.URLError as exc:
        print(f"API 请求失败: {exc.reason}")
        print("请确认 API 服务已启动: uvicorn lyric_dedup_server.app:app --host 0.0.0.0 --port 8000")
        sys.exit(1)


def _call_rhythm_check(url: str, top_k: int, api_base: str) -> dict:
    """调用曲去重查询 API"""
    payload = json.dumps({"url": url, "top_k": top_k}).encode("utf-8")
    req = urllib.request.Request(
        f"{api_base.rstrip('/')}/api/v1/rhythm/check",
        data=payload,
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=120) as resp:
            return json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        error_body = exc.read().decode("utf-8", errors="replace")
        print(f"API 请求失败 (HTTP {exc.code}): {error_body}")
        sys.exit(1)
    except urllib.error.URLError as exc:
        print(f"API 请求失败: {exc.reason}")
        print("请确认 API 服务已启动: uvicorn lyric_dedup_server.app:app --host 0.0.0.0 --port 8000")
        sys.exit(1)


if __name__ == "__main__":
    main()