Create-knowledge-find-decode/scripts/classify_creation_items.py

#!/usr/bin/env python3
"""补跑 430 query 采集结果的创作知识判断。

示例：
  CK_ENV_FILE=.env PYTHONPATH=. .venv/bin/python scripts/classify_creation_items.py \
    --run-id creation-search-full-20260629 --platform xiaohongshu,weixin --workers 4

  CK_ENV_FILE=.env PYTHONPATH=. .venv/bin/python scripts/classify_creation_items.py \
    --run-id creation-search-full-20260629 --platform xiaohongshu --provider qwen --model qwen3.7-plus --workers 4

  CK_ENV_FILE=.env PYTHONPATH=. .venv/bin/python scripts/classify_creation_items.py \
    --run-id creation-search-full-20260629 --platform weixin --provider ark --model ep-20260506151915-jqvw7 --workers 4
"""
from __future__ import annotations

import argparse
import concurrent.futures as cf
import os
import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from acquisition import store
from acquisition.classify import classify_imgtext, classify_video
from acquisition.creation_search import PLATFORMS, prompt_version
from core.config import Settings


def _parse_platforms(value: str | None) -> list[str] | None:
    if not value:
        return None
    out = [p.strip() for p in value.split(",") if p.strip()]
    bad = [p for p in out if p not in PLATFORMS]
    if bad:
        raise SystemExit(f"unsupported platform(s): {', '.join(bad)}")
    return out


def _classify_one(item: dict, settings: Settings) -> tuple[int, int | None, str, str, str, str]:
    platform = item["platform"]
    if platform == "douyin":
        version = prompt_version("classify_video")
        is_creation, reason, knowledge, _points = classify_video(
            {
                "platform": platform,
                "title": item.get("title") or "",
                "body_text": item.get("body_text") or "",
                "video": item.get("video_url") or "",
            },
            settings,
        )
    else:
        version = prompt_version("classify_imgtext")
        is_creation, reason, knowledge, _points = classify_imgtext(
            {
                "platform": platform,
                "title": item.get("title") or "",
                "body_text": item.get("body_text") or "",
                "images": item.get("image_urls") or [],
            },
            settings,
        )
    return item["id"], is_creation, reason, knowledge, version, reason if is_creation is None else ""


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--run-id", help="默认使用最新 creation_search_run")
    ap.add_argument("--platform", help="平台，多个用逗号分隔；默认全部")
    ap.add_argument("--limit", type=int, default=0, help="最多补判多少条；0=不限制")
    ap.add_argument("--workers", type=int, default=3)
    ap.add_argument("--only-missing", action="store_true", help="只补完全未分类，不重试失败分类")
    ap.add_argument("--provider", choices=["auto", "openrouter", "ark", "qwen", "bailian"], help="本次补判使用的模型通道")
    ap.add_argument("--model", help="本次补判使用的模型名/方舟接入点，如 qwen3.7-plus 或 ep-...")
    args = ap.parse_args()

    if args.provider:
        os.environ["CLASSIFY_PROVIDER"] = args.provider
    if args.model:
        os.environ["CLASSIFY_MODEL"] = args.model

    settings = Settings.from_env(os.getenv("CK_ENV_FILE", ".env"))
    platforms = _parse_platforms(args.platform)
    conn = store.connect()
    rows = store.creation_items_to_classify(
        conn,
        run_id=args.run_id,
        platforms=platforms,
        retry_failed=not args.only_missing,
        limit=args.limit if args.limit > 0 else None,
    )
    print(f"待补判 {len(rows)} 条 platforms={','.join(platforms or PLATFORMS)} workers={args.workers}")
    if not rows:
        conn.close()
        return

    done = fail = 0
    with cf.ThreadPoolExecutor(max_workers=max(1, args.workers)) as ex:
        futs = {ex.submit(_classify_one, item, settings): item for item in rows}
        for fut in cf.as_completed(futs):
            item = futs[fut]
            try:
                item_id, is_creation, reason, knowledge, version, error = fut.result()
            except Exception as exc:
                item_id = item["id"]
                is_creation = None
                reason = f"判定失败: {str(exc)[:80]}"
                knowledge = ""
                version = ""
                error = reason
            store.upsert_creation_classification(
                conn,
                item_id,
                is_creation,
                reason=reason,
                knowledge=knowledge,
                prompt_version=version,
                error=error,
                ts=int(time.time()),
            )
            done += 1
            if is_creation is None:
                fail += 1
            if done % 20 == 0 or done == len(rows):
                print(f"  {done}/{len(rows)} done failed={fail}", flush=True)
    conn.close()
    print(f"完成：{done} 条，失败 {fail} 条")


if __name__ == "__main__":
    main()