Agent/tests/test_liblibai_tool.py

"""测试 liblibai_controlnet 工具 — 通过 Router API 调用

用法：
    uv run python tests/test_liblibai_tool.py                    # 运行测试
    uv run python tests/test_liblibai_tool.py --search           # 仅搜索工具
"""

import argparse
import json
import os
import sys
from typing import Any

if sys.platform == "win32":
    sys.stdout.reconfigure(encoding="utf-8")

import httpx

BASE_URL = os.environ.get("TOOL_AGENT_ROUTER_URL", "http://127.0.0.1:8001")
DEFAULT_TIMEOUT = 30.0
CALL_TIMEOUT = 300.0

IMAGE_URL = "https://liblibai-airship-temp.oss-cn-beijing.aliyuncs.com/aliyun-cn-prod/73ed6ae42b144d21bf566e05b5a6c138.png"


def check_connection():
    try:
        httpx.get(f"{BASE_URL}/health", timeout=3)
    except httpx.ConnectError:
        print(f"ERROR: Cannot connect to {BASE_URL}")
        print("Please start the service first:")
        print("  uv run python -m tool_agent")
        sys.exit(1)


def _run_tool(
    tool_id: str, params: dict[str, Any], timeout: float = CALL_TIMEOUT
) -> tuple[bool, str | None, Any]:
    """POST /run_tool。成功返回 (True, None, result)；失败 (False, message, None)。"""
    resp = httpx.post(
        f"{BASE_URL}/run_tool",
        json={"tool_id": tool_id, "params": params},
        timeout=timeout,
    )
    print(f"  Status : {resp.status_code}")
    if resp.status_code != 200:
        return False, f"HTTP {resp.status_code}: {resp.text[:300]}", None
    try:
        data = resp.json()
    except Exception as e:
        return False, f"Invalid JSON: {e}", None
    if data.get("status") != "success":
        return False, data.get("error") or str(data), None
    result = data.get("result")
    if isinstance(result, dict) and result.get("status") == "error":
        return False, str(result.get("error", result)), None
    return True, None, result


def test_health():
    print("=== Health Check ===")
    resp = httpx.get(f"{BASE_URL}/health", timeout=DEFAULT_TIMEOUT)
    print(f"  Status : {resp.status_code}")
    print(f"  Body   : {json.dumps(resp.json(), ensure_ascii=False, indent=4)}")
    assert resp.status_code == 200
    print("  [PASS]")


def test_search_tools():
    print("=== Search Tools (keyword='liblib') ===")
    resp = httpx.get(f"{BASE_URL}/tools", timeout=DEFAULT_TIMEOUT)
    if resp.status_code != 200:
        print(f"  ERROR: GET /tools failed: {resp.status_code}")
        print("  [FAIL]")
        return
    data = resp.json()
    tools = [
        t for t in data.get("tools", [])
        if "liblib" in t.get("tool_id", "").lower()
        or "liblib" in t.get("name", "").lower()
    ]
    print(f"  Found  : {len(tools)}")
    for t in tools:
        print(f"\n  [{t['tool_id']}]")
        print(f"    name    : {t['name']}")
        print(f"    state   : {t['state']}")
        print(f"    category: {t.get('category', '')}")
        print(f"    runtime : {t.get('backend_runtime', '')}")
        props = t.get("input_schema", {}).get("properties", {})
        required = t.get("input_schema", {}).get("required", [])
        print(f"    params ({len(props)}):")
        for name, info in props.items():
            req = "*" if name in required else " "
            desc = info.get("description", "")
            default_str = f"  default={info['default']}" if info.get("default") is not None else ""
            print(f"      {req} {name:<20} {info.get('type','any'):<10} {desc}{default_str}")
    print("\n  [PASS]")


def test_controlnet():
    """POST /run_tool → liblibai_controlnet"""
    print("=== Test liblibai_controlnet ===")

    # 确认工具已注册
    tid = "liblibai_controlnet"
    try:
        tr = httpx.get(f"{BASE_URL}/tools", timeout=DEFAULT_TIMEOUT)
        tr.raise_for_status()
        ids = {t["tool_id"] for t in tr.json().get("tools", [])}
        if tid not in ids:
            print(f"  ERROR : 注册表中无 {tid!r}")
            print("  [FAIL]")
            return
        print(f"  tool_id: {tid} (已注册)")
    except Exception as e:
        print(f"  ERROR : GET /tools 失败: {e}")
        print("  [FAIL]")
        return

    prompt = os.environ.get(
        "LIBLIBAI_TEST_PROMPT",
        "simple white line art, cat, black background",
    )
    params: dict[str, Any] = {
        "image": IMAGE_URL,
        "prompt": prompt,
        "negative_prompt": "lowres, bad anatomy, text, error",
        "width": 512,
        "height": 512,
        "steps": 20,
        "cfg_scale": 7,
        "img_count": 1,
        "control_weight": 1.0,
        "preprocessor": 1,
        "canny_low": 100,
        "canny_high": 200,
    }

    print(f"  image  : {IMAGE_URL}")
    print(f"  prompt : {prompt}")
    print(f"  calling {tid} ...")

    try:
        ok, err, result = _run_tool(tid, params, timeout=CALL_TIMEOUT)
        if not ok:
            print(f"  ERROR : {err}")
            print("  [FAIL]")
            return
        if not isinstance(result, dict):
            print(f"  ERROR : 非 dict 结果: {type(result)}")
            print("  [FAIL]")
            return

        print(f"  Result :")
        print(f"    status : {result.get('status')}")
        print(f"    task_id: {result.get('task_id')}")
        images = result.get("images", [])
        if images:
            print(f"    images : {len(images)}")
            for i, url in enumerate(images):
                print(f"      [{i+1}] {url}")
            print("  [PASS]")
        else:
            print("    WARN : 未返回图片，任务可能仍在处理中")
            print(f"    keys : {list(result.keys())}")
            print(f"    截断 : {str(result)[:400]}...")
            print("  [PASS] (submitted)")
    except httpx.TimeoutException:
        print(f"  ERROR : Request timeout ({CALL_TIMEOUT:.0f}s)")
        print("  [FAIL]")
    except Exception as e:
        print(f"  ERROR : {e}")
        print("  [FAIL]")


def main():
    parser = argparse.ArgumentParser(description="LibLibAI ControlNet Tool Test")
    parser.add_argument("--search", action="store_true",
                        help="only search and list liblib tools")
    args = parser.parse_args()

    print(f"Target: {BASE_URL}\n")
    check_connection()
    test_health()

    if args.search:
        print()
        test_search_tools()
    else:
        print()
        test_search_tools()
        print()
        test_controlnet()

    print("\n=== DONE ===")


if __name__ == "__main__":
    main()