holy-python/skills/holy-crab/scripts/query_tasks.py

#!/usr/bin/env python3
"""
query_tasks.py — 列出 Holy Crab 所有任务
用法: python3 query_tasks.py [任务ID 关键词]

Examples:
    python3 query_tasks.py
    python3 query_tasks.py HC-20260727-ABC12
"""

import json
import os
import sys
from pathlib import Path

SCRIPT_DIR = Path(__file__).resolve().parent
DEFAULT_BACKEND = "http://localhost:8000"


def load_cookie() -> str | None:
    # 优先级：环境变量 > ~/.holy_crab_env > ~/.holy_crab_cookie
    env = os.environ.get("HC_SESSION")
    if env:
        return env
    env_file = Path.home() / ".holy_crab_env"
    if env_file.exists():
        for line in env_file.read_text().strip().splitlines():
            line = line.strip()
            if line.startswith("HC_SESSION=") or line.startswith("hc_session="):
                val = line.split("=", 1)[1].strip().strip('"').strip("'")
                if val:
                    return val
    cookie_file = Path.home() / ".holy_crab_cookie"
    if cookie_file.exists():
        return cookie_file.read_text().strip()
    return None


def api_get(path: str, params: dict | None = None) -> dict:
    backend = os.environ.get("HC_BACKEND_URL", DEFAULT_BACKEND).rstrip("/")
    url = f"{backend}/api/v1{path}"
    cookie = load_cookie()
    headers = {}
    if cookie:
        headers["Cookie"] = f"hc_session={cookie}"

    import urllib.request
    req = urllib.request.Request(url, headers=headers)
    if params:
        import urllib.parse
        req.full_url = f"{url}?{urllib.parse.urlencode(params)}"
    try:
        with urllib.request.urlopen(req, timeout=15) as resp:
            return json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as exc:
        body = exc.read().decode("utf-8", errors="replace")
        try:
            err = json.loads(body)
            print(f"❌ API错误 {exc.code}: {err}", file=sys.stderr)
        except Exception:
            print(f"❌ HTTP {exc.code}: {body[:300]}", file=sys.stderr)
        sys.exit(1)


def format_status(status: str) -> str:
    mapping = {
        "completed": "✅ 完成",
        "failed": "❌ 失败",
        "cancelled": "🚫 取消",
        "partial": "⚠️ 部分",
        "running": "⏳ 进行中",
        "waiting": "⏳ 等待中",
    }
    return mapping.get(status, status)


def main() -> None:
    task_id_filter = sys.argv[1] if len(sys.argv) > 1 else None

    if task_id_filter:
        # 精确查单个任务
        data = api_get(f"/tasks/{task_id_filter}")
        task = data.get("data", {})
        print(f"\n{'─'*60}")
        print(f"  任务ID    : {task.get('task_id')}")
        print(f"  名称      : {task.get('name')}")
        print(f"  关键词    : {', '.join(task.get('keywords', []))}")
        print(f"  状态      : {format_status(task.get('status'))}")
        print(f"  进度      : {task.get('progress_percent')}%")
        print(f"  模式      : {task.get('mode')}")
        print(f"  创建时间  : {task.get('created_at')}")
        print(f"  完成时间  : {task.get('completed_at') or '—'}")
        if task.get("error"):
            print(f"  错误      : [{task['error'].get('code')}] {task['error'].get('message')}")
        # 检查 extracted
        result = task.get("result", {})
        extracted = result.get("extracted")
        if extracted:
            meta = extracted.get("提取元数据", {})
            stats = meta.get("召回统计", {})
            print(f"\n  📊 数据清洗结果:")
            print(f"     schema    : {extracted.get('schema_version')}")
            print(f"     标签总数  : {stats.get('标签总数', 'N/A')}")
            print(f"     召回总数  : {stats.get('实际召回内容总数', 'N/A')}")
            print(f"     榜单数量  : {len(extracted.get('榜单', []))}")
            print(f"     产品数量  : {len(extracted.get('产品详情', []))}")
            src = extracted.get("参考来源笔记总量", {})
            print(f"     笔记总量  : {src.get('原始文本', 'N/A')}")
        else:
            print(f"\n  ⏳ 任务尚未完成或无 extracted 数据")
        print(f"{'─'*60}\n")
    else:
        # 列表
        data = api_get("/tasks", {"page_size": 50})
        items = data.get("data", {}).get("items", [])
        total = data.get("data", {}).get("total", 0)
        print(f"\n{'─'*60}")
        print(f"  Holy Crab 任务列表  (共 {total} 个)")
        print(f"{'─'*60}")
        print(f"  {'任务ID':<22} {'状态':<14} {'关键词':<16} {'创建时间':<12} {'文件'}")
        print(f"  {'─'*60}")
        for t in items:
            keywords = ",".join(t.get("keywords", [])[:2])
            if len(",".join(t.get("keywords", []))) > 16:
                keywords += "..."
            status = format_status(t.get("status", ""))
            files = len(t.get("files", []))
            ts = (t.get("created_at") or "")[:10]
            print(f"  {t.get('task_id'):<22} {status:<14} {keywords:<16} {ts:<12} {files}f")
        print(f"{'─'*60}\n")


if __name__ == "__main__":
    main()