Tool_Agent/tests/test_knowhub_modify.py

"""批量给 KnowHub 工具表添加 toolhub_items 字段

规则：
- 已接入的工具：填入对应的 Tool Agent 工具 ID 和名称
- 未接入的工具：只加空的 toolhub_items 字段
"""

import httpx
import os
import json

KNOWHUB_API = os.getenv("KNOWHUB_API", "http://43.106.118.91:9999")
TIMEOUT = 60.0

# Tool Agent registry 中已接入的工具 → KnowHub 工具表的映射
# key: KnowHub tool_id, value: toolhub_items 列表
TOOLHUB_MAPPING = {
    "tools/image_gen/comfyui": [
        {"launch_comfy_env": "启动云端 ComfyUI Docker 环境，返回 server_id 和 comfy_url"},
        {"runcomfy_workflow_executor": "在已就绪的 RunComfy 机器上提交 ComfyUI 工作流，上传输入文件，监听执行状态，下载结果图片"},
        {"runcomfy_stop_env": "停止并删除 RunComfy 机器实例，释放资源"},
    ],
}


def get_all_tools() -> list[dict]:
    """获取 KnowHub 所有工具"""
    resp = httpx.get(f"{KNOWHUB_API}/api/resource", params={"limit": 1000}, timeout=TIMEOUT)
    resp.raise_for_status()
    data = resp.json()
    results = data.get("results", [])
    return [r for r in results if isinstance(r, dict) and r.get("id", "").startswith("tools/")]


def patch_tool(tool_id: str, metadata: dict) -> bool:
    """PATCH 更新工具 metadata"""
    resp = httpx.patch(
        f"{KNOWHUB_API}/api/resource/{tool_id}",
        json={"metadata": metadata},
        timeout=TIMEOUT
    )
    resp.raise_for_status()
    return True


def main():
    print(f"KnowHub API: {KNOWHUB_API}")
    print("=" * 60)

    tools = get_all_tools()
    print(f"共 {len(tools)} 个工具\n")

    updated = 0
    skipped = 0
    failed = 0

    for tool in tools:
        tool_id = tool["id"]
        metadata = tool.get("metadata", {})

        # 已经有 toolhub_items 且非空的跳过
        if metadata.get("toolhub_items"):
            print(f"  SKIP  {tool_id} (已有 toolhub_items)")
            skipped += 1
            continue

        # 查映射表，有则填入，无则空列表
        items = TOOLHUB_MAPPING.get(tool_id, [])
        metadata["toolhub_items"] = items

        try:
            patch_tool(tool_id, metadata)
            tag = f"{len(items)} items" if items else "empty"
            print(f"  OK    {tool_id} ({tag})")
            updated += 1
        except Exception as e:
            print(f"  FAIL  {tool_id}: {e}")
            failed += 1

    print(f"\n{'=' * 60}")
    print(f"完成: 更新 {updated}, 跳过 {skipped}, 失败 {failed}")


if __name__ == "__main__":
    main()