Tool_Agent/tests/test_liblibai_tool.py

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

用法：
    1. 先启动 Router：uv run python -m tool_agent
    2. 运行测试：python tests/test_liblibai_tool.py
"""

import sys
import time

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

import httpx

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


def main():
    print("=" * 50)
    print("测试 liblibai_controlnet 工具")
    print("=" * 50)

    # 1. 检查 Router 是否在线
    try:
        resp = httpx.get(f"{ROUTER_URL}/health", timeout=3)
        print(f"Router 状态: {resp.json()}")
    except httpx.ConnectError:
        print(f"无法连接 Router ({ROUTER_URL})")
        print("请先启动: uv run python -m tool_agent")
        sys.exit(1)

    # 2. 搜索工具，确认已注册
    print("\n--- 获取工具 ---")
    resp = httpx.get(f"{ROUTER_URL}/tools")
    tools = resp.json()
    liblib_tools = [t for t in tools["tools"] if "liblib" in t["tool_id"] or "liblib" in t["name"].lower()]
    print(f"找到 {len(liblib_tools)} 个匹配 liblib 的工具")
    for t in liblib_tools:
        print(f"  {t['tool_id']}: {t['name']} (state={t['state']})")

    # 3. 调用 liblibai_controlnet 工具
    print("\n--- 调用 ControlNet 生图 ---")
    print(f"图片: {IMAGE_URL}")
    print(f"提示词: simple white line art, cat, black background")
    print("提交中...")

    resp = httpx.post(
        f"{ROUTER_URL}/run_tool",
        json={
            "tool_id": "liblibai_controlnet",
            "params": {
                "image": IMAGE_URL,
                "prompt": "simple white line art, cat, black background",
                "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
            }
        },
        timeout=300  # 生图可能需要几分钟
    )

    result = resp.json()
    print(f"\n响应状态: {result.get('status')}")

    if result.get("status") == "success":
        data = result.get("result", {})
        print(f"任务状态: {data.get('status')}")
        print(f"任务 ID: {data.get('task_id')}")
        images = data.get("images", [])
        if images:
            print(f"生成图片数: {len(images)}")
            for i, url in enumerate(images):
                print(f"  [{i+1}] {url}")
            print("\n测试通过!")
        else:
            print("未返回图片，可能任务仍在处理中")
    else:
        print(f"错误: {result.get('error')}")


if __name__ == "__main__":
    main()