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

#!/usr/bin/env python3
"""真实运行 430 query 搜索/详情/媒体/分类。

示例：
  PYTHONPATH=. CK_ENV_FILE=.env python scripts/run_creation_search.py \
    --query "公共安全 视频 灵感 怎么做" --search-limit 10 --display-limit 5
"""
from __future__ import annotations

import argparse
import concurrent.futures as cf
import datetime as dt
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.creation_search import PLATFORMS, load_creation_queries, run_platform_query
from core.config import Settings


def _run_id() -> str:
    return "creation-search-" + dt.datetime.now().strftime("%Y%m%d-%H%M%S")


def _parse_platforms(value: str | None) -> list[str]:
    if not value:
        return list(PLATFORMS)
    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 run_platform_workers(*, run_id: str, queries: list[str], platforms: list[str],
                         settings: Settings, search_limit: int, display_limit: int,
                         classify: bool, skip_done: bool,
                         db_path=store.DB_PATH) -> int:
    """Run long-lived platform workers. Return failed platform-job count."""

    def _platform_worker(platform: str) -> tuple[str, int, int]:
        c = store.connect(db_path)
        ok = fail = 0
        current_q = ""
        try:
            for i, q in enumerate(queries, start=1):
                current_q = q
                if skip_done and store.creation_job_is_done(
                    c, run_id, q, platform, display_limit=display_limit
                ):
                    ok += 1
                    print(f"[{platform}] {i}/{len(queries)} skip done {q}", flush=True)
                    continue
                print(f"[{platform}] {i}/{len(queries)} {q}", flush=True)
                res = run_platform_query(
                    c,
                    run_id=run_id,
                    query=q,
                    platform=platform,
                    settings=settings,
                    search_limit=search_limit,
                    display_limit=display_limit,
                    classify=classify,
                )
                if res["status"] in ("done", "partial"):
                    ok += 1
                else:
                    fail += 1
                print(
                    f"[{platform}] {res['status']} display={res['display_count']} "
                    f"{res.get('error') or ''}",
                    flush=True,
                )
        except Exception as exc:
            fail += 1
            msg = f"worker异常: {str(exc)[:160]}"
            if current_q:
                store.update_creation_job(
                    c, run_id, current_q, platform, status="failed",
                    error=msg, ts=int(time.time()),
                )
            print(f"[{platform}] {msg}", flush=True)
        finally:
            c.close()
        return platform, ok, fail

    failed = 0
    with cf.ThreadPoolExecutor(max_workers=len(platforms)) as ex:
        futs = [ex.submit(_platform_worker, p) for p in platforms]
        for fut in cf.as_completed(futs):
            platform, ok, fail = fut.result()
            failed += fail
            print(f"[{platform}] finished ok={ok} failed={fail}")
    return failed


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--run-id", default=_run_id())
    ap.add_argument("--query", help="只跑这一条 query；不传则读取 creation_demo.json")
    ap.add_argument("--limit-queries", type=int, default=0, help="从 query 文件取前 N 条；0=不截断")
    ap.add_argument("--platform", help="只跑指定平台；多个用逗号分隔")
    ap.add_argument("--search-limit", type=int, default=10)
    ap.add_argument("--display-limit", type=int, default=5)
    ap.add_argument("--query-file", default=str(Path("data/queries/creation_demo.json")))
    ap.add_argument("--no-classify", action="store_true", help="只采集不调用 AI 判断")
    ap.add_argument("--resume", action="store_true", help="保留同 run_id 已有结果，继续未完成 job")
    ap.add_argument("--skip-done", action="store_true", help="跳过已 done 且 display_count>=display_limit 的 job")
    args = ap.parse_args()

    settings = Settings.from_env(os.getenv("CK_ENV_FILE", ".env"))
    queries = [args.query.strip()] if args.query else load_creation_queries(args.query_file)
    if args.limit_queries and args.limit_queries > 0:
        queries = queries[:args.limit_queries]
    platforms = _parse_platforms(args.platform)
    if not queries:
        raise SystemExit("no queries to run")

    conn = store.connect()
    if args.resume and store.creation_run_exists(conn, args.run_id):
        print(f"resume run_id={args.run_id}")
    else:
        store.create_creation_run(
            conn,
            args.run_id,
            total_queries=len(queries),
            note=f"platforms={','.join(platforms)} search_limit={args.search_limit} display_limit={args.display_limit}",
            ts=int(time.time()),
        )
    for q in queries:
        for p in platforms:
            store.ensure_creation_job(conn, args.run_id, q, p, ts=int(time.time()))
    conn.close()

    print(f"run_id={args.run_id} queries={len(queries)} platforms={','.join(platforms)}")
    print("平台接口限速：同平台搜索/详情共享 10-12s 间隔；OSS/本地下载/AI 判断不占平台闸。")

    failed = run_platform_workers(
        run_id=args.run_id,
        queries=queries,
        platforms=platforms,
        settings=settings,
        search_limit=args.search_limit,
        display_limit=args.display_limit,
        classify=not args.no_classify,
        skip_done=args.skip_done or args.resume,
    )

    c = store.connect()
    store.finish_creation_run(c, args.run_id, status="finished" if failed == 0 else "partial", ts=int(time.time()))
    c.close()
    print(f"done run_id={args.run_id} failed_platform_jobs={failed}")


if __name__ == "__main__":
    main()