Tool_Agent/tests/test_knowhub_query.py

"""测试 KnowHub 工具表查询功能"""

import httpx
import os
from typing import Any

# KnowHub API 地址
KNOWHUB_API = os.getenv("KNOWHUB_API", "http://43.106.118.91:9999")


def test_list_all_resources():
    """测试：列出所有资源"""
    print("\n=== 测试：列出所有资源 ===")
    try:
        resp = httpx.get(f"{KNOWHUB_API}/api/resource", params={"limit": 1000}, timeout=60.0)
        resp.raise_for_status()
        data = resp.json()

        print(f"状态码: {resp.status_code}")
        print(f"返回数据类型: {type(data)}")

        # 检查返回格式
        if isinstance(data, dict):
            print(f"返回字典的键: {data.keys()}")
            # API 返回格式: {"results": [...], "count": N}
            if "results" in data:
                data = data["results"]
            elif "data" in data:
                data = data["data"]

        if not isinstance(data, list):
            print(f"⚠️ 返回数据不是列表")
            return []

        print(f"总资源数: {len(data)}")

        # 过滤出工具
        tools = [r for r in data if isinstance(r, dict) and r.get("id", "").startswith("tools/")]
        print(f"工具数量: {len(tools)}")

        if tools:
            print("\n前 3 个工具:")
            for tool in tools[:3]:
                print(f"  - {tool['id']}: {tool.get('title', 'N/A')}")

        return tools
    except Exception as e:
        print(f"❌ 错误: {e}")
        import traceback
        traceback.print_exc()
        return []


def test_get_tool_detail(tool_id: str):
    """测试：获取工具详情"""
    print(f"\n=== 测试：获取工具详情 ({tool_id}) ===")
    try:
        resp = httpx.get(f"{KNOWHUB_API}/api/resource/{tool_id}", timeout=10.0)
        resp.raise_for_status()
        tool = resp.json()

        print(f"状态码: {resp.status_code}")
        print(f"ID: {tool['id']}")
        print(f"标题: {tool.get('title', 'N/A')}")

        metadata = tool.get("metadata", {})
        print(f"\nMetadata:")
        for key, value in metadata.items():
            # 长列表截断显示
            if isinstance(value, list) and len(value) > 3:
                print(f"  - {key}: [{len(value)} items] {value[:3]}...")
            else:
                print(f"  - {key}: {value}")

        return tool
    except Exception as e:
        print(f"❌ 错误: {e}")
        return None


def test_search_tools_by_category(category: str):
    """测试：按分类搜索工具"""
    print(f"\n=== 测试：按分类搜索工具 (category={category}) ===")
    try:
        resp = httpx.get(f"{KNOWHUB_API}/api/resource", params={"limit": 1000}, timeout=60.0)
        resp.raise_for_status()
        data = resp.json()

        # 检查返回格式
        if isinstance(data, dict):
            # API 返回格式: {"results": [...], "count": N}
            if "results" in data:
                data = data["results"]
            elif "data" in data:
                data = data["data"]

        if not isinstance(data, list):
            print(f"⚠️ 返回数据不是列表")
            return []

        # 客户端过滤
        tools = [
            r for r in data
            if isinstance(r, dict) and r.get("id", "").startswith("tools/")
            and r.get("metadata", {}).get("category") == category
        ]

        print(f"找到 {len(tools)} 个 {category} 类工具:")
        for tool in tools[:5]:
            print(f"  - {tool['id']}: {tool.get('title', 'N/A')}")

        return tools
    except Exception as e:
        print(f"❌ 错误: {e}")
        import traceback
        traceback.print_exc()
        return []


def test_search_knowledge_by_tool(query: str, top_k: int = 5):
    """测试：搜索工具相关知识"""
    print(f"\n=== 测试：搜索工具相关知识 (query={query}) ===")
    try:
        resp = httpx.get(
            f"{KNOWHUB_API}/api/knowledge/search",
            params={"q": query, "top_k": top_k, "min_score": 3},
            timeout=60.0
        )
        resp.raise_for_status()
        results = resp.json()

        print(f"状态码: {resp.status_code}")
        print(f"返回数据类型: {type(results)}")

        # 检查返回格式
        if isinstance(results, dict):
            print(f"返回字典的键: {results.keys()}")
            if "results" in results:
                results = results["results"]
            elif "data" in results:
                results = results["data"]

        if not isinstance(results, list):
            print(f"⚠️ 返回数据不是列表")
            return []

        print(f"找到 {len(results)} 条知识:")

        for i, item in enumerate(results[:3], 1):
            if isinstance(item, dict):
                print(f"\n  {i}. {item.get('id', 'N/A')}")
                title = item.get('title', 'N/A')
                print(f"     标题: {title[:50] if isinstance(title, str) else title}")
                print(f"     分数: {item.get('score', 0):.2f}")
                print(f"     关联工具: {item.get('resource_ids', [])}")

        return results
    except Exception as e:
        print(f"❌ 错误: {e}")
        import traceback
        traceback.print_exc()
        return []


def test_list_tool_knowledge(tags: str = "tool"):
    """测试：列出所有工具相关知识"""
    print(f"\n=== 测试：列出所有工具相关知识 (tags={tags}) ===")
    try:
        resp = httpx.get(
            f"{KNOWHUB_API}/api/knowledge",
            params={"tags": tags, "status": "approved,checked", "page_size": 200},
            timeout=60.0
        )
        resp.raise_for_status()
        results = resp.json()

        print(f"状态码: {resp.status_code}")
        print(f"返回数据类型: {type(results)}")

        # 检查返回格式
        if isinstance(results, dict):
            print(f"返回字典的键: {results.keys()}")
            if "results" in results:
                results = results["results"]
            elif "data" in results:
                results = results["data"]
            elif "items" in results:
                results = results["items"]

        if not isinstance(results, list):
            print(f"⚠️ 返回数据不是列表")
            return []

        print(f"找到 {len(results)} 条工具知识")

        if results:
            print("\n前 3 条:")
            for item in results[:3]:
                if isinstance(item, dict):
                    title = item.get('title', 'N/A')
                    print(f"  - {item.get('id', 'N/A')}: {title[:40] if isinstance(title, str) else title}")
                    print(f"    关联工具: {item.get('resource_ids', [])}")

        return results
    except Exception as e:
        print(f"❌ 错误: {e}")
        import traceback
        traceback.print_exc()
        return []


def main():
    """运行所有测试"""
    print(f"KnowHub API: {KNOWHUB_API}")
    print("=" * 60)

    # 1. 列出所有资源
    tools = test_list_all_resources()

    # 2. 获取 ComfyUI 工具详情（验证 toolhub_items）
    test_get_tool_detail("tools/image_gen/comfyui")

    # 3. 获取第一个工具的详情
    if tools:
        first_tool_id = tools[0]["id"]
        if first_tool_id != "tools/image_gen/comfyui":
            test_get_tool_detail(first_tool_id)

    # 3. 按分类搜索
    test_search_tools_by_category("plugin")

    # 4. 搜索工具相关知识
    test_search_knowledge_by_tool("IP-Adapter")

    # 5. 列出所有工具知识
    test_list_tool_knowledge()

    print("\n" + "=" * 60)
    print("测试完成！")


if __name__ == "__main__":
    main()