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

"""用生成的 query 真实搜帖子/视频（抖音 + 微信公众号 + 小红书），媒体存本地 data/。

每法取前 N 条 query，每条 query 每个渠道取 top-K：
  抖音：search(top-1) → detail → 下载视频+封面（搜索/详情各自限速）
  微信公众号：search(top-K) → 下载每个封面（封面在搜索回包里，无需 detail）
  小红书：search(top-K) → 下载每个封面（封面在 note_card 里，无需 detail；视频帖无直链）
限流：三个平台的「搜索」各自独立闸，两次搜索间隔随机 10~12s；抖音「详情」走更轻的 3~5s 闸。
产出每条记录：{method, query, douyin/weixin/xiaohongshu:{ok:[...],error}}。
用法：PYTHONPATH=. python scripts/run_search.py
"""
from __future__ import annotations

import hashlib
import json
from pathlib import Path

from acquisition.crawler import RateLimiter, fetch_post_detail
from acquisition.search import search_keyword, search_weixin, search_xiaohongshu
from core.config import Settings
from creation_knowledge.integrations.video_extract import _default_download

ROOT = Path(__file__).resolve().parent.parent
DATA = ROOT / "data"
DEMO = DATA / "queries" / "demo.json"
OUT = DATA / "queries" / "search_results.json"
DOUYIN_ON = False                       # 抖音限流暂停：False 时不发抖音请求，沿用上次结果里的抖音
N = 6                                    # 每法取前 N 条 query
K_DOUYIN = 1                             # 抖音每条 query 只取第 1 个视频（躲限流）
K_WEIXIN = 5                             # 微信公众号每条 query 取前 K 个封面
K_XHS = 5                               # 小红书每条 query 取前 K 个封面
SEARCH_MIN, SEARCH_MAX = 10.0, 12.0     # 每平台两次搜索之间随机间隔（秒）


def _dl(url: str, platform: str, dst: Path, public: str):
    """下载到 dst，返回公开路径；失败返回 None。"""
    try:
        dst.parent.mkdir(parents=True, exist_ok=True)
        dst.write_bytes(_default_download(url, platform))
        return public
    except Exception:
        return None


def douyin_topk(query: str, settings: Settings, search_rl: RateLimiter, detail_rl: RateLimiter) -> dict:
    """抖音搜 top-K 视频：返回 {ok:[{title,url,cover,video}], error}。搜索/详情各走各的闸。"""
    try:
        ids = search_keyword(query, platform="douyin", content_type="视频", limit=K_DOUYIN,
                             settings=settings, rate_limiter=search_rl)
    except Exception as exc:
        return {"ok": [], "error": f"搜索失败: {str(exc)[:50]}"}
    if not ids:
        return {"ok": [], "error": "未搜到"}
    recs = []
    for vid in ids[:K_DOUYIN]:
        try:
            post = fetch_post_detail(vid, settings=settings, rate_limiter=detail_rl)
        except Exception:
            continue
        base, pub = DATA / "search" / "douyin" / post.id, f"/data/search/douyin/{post.id}"
        rec = {"title": post.title, "url": post.url, "cover": None, "video": None}
        if post.image_urls:
            rec["cover"] = _dl(post.image_urls[0], "douyin", base / "cover.jpg", pub + "/cover.jpg")
        if post.video_urls:
            rec["video"] = _dl(post.video_urls[0], "douyin", base / "video.mp4", pub + "/video.mp4")
        if rec["video"] or rec["cover"]:
            recs.append(rec)
    return {"ok": recs, "error": None if recs else "详情/下载均失败"}


def weixin_topk(query: str, settings: Settings, rl: RateLimiter) -> dict:
    """微信公众号搜 top-K 文章封面：返回 {ok:[{title,url,nick,cover}], error}。"""
    try:
        arts = search_weixin(query, limit=K_WEIXIN, settings=settings, rate_limiter=rl)
    except Exception as exc:
        return {"ok": [], "error": f"搜索失败: {str(exc)[:50]}"}
    if not arts:
        return {"ok": [], "error": "未搜到"}
    recs = []
    for a in arts[:K_WEIXIN]:
        h = hashlib.md5(a["url"].encode()).hexdigest()[:16]
        base, pub = DATA / "search" / "weixin" / h, f"/data/search/weixin/{h}"
        rec = {"title": a["title"], "url": a["url"], "nick": a["nick_name"], "cover": None}
        if a["cover_url"]:
            rec["cover"] = _dl(a["cover_url"], "weixin", base / "cover.jpg", pub + "/cover.jpg")
        recs.append(rec)
    return {"ok": recs, "error": None}


def xiaohongshu_topk(query: str, settings: Settings, rl: RateLimiter) -> dict:
    """小红书搜 top-K 帖子封面：返回 {ok:[{title,url,nick,cover}], error}。封面来自搜索回包，无需 detail。"""
    try:
        posts = search_xiaohongshu(query, content_type="图文", limit=K_XHS,
                                   settings=settings, rate_limiter=rl)
    except Exception as exc:
        return {"ok": [], "error": f"搜索失败: {str(exc)[:50]}"}
    if not posts:
        return {"ok": [], "error": "未搜到"}
    recs = []
    for p in posts[:K_XHS]:
        base, pub = DATA / "search" / "xiaohongshu" / p["id"], f"/data/search/xiaohongshu/{p['id']}"
        rec = {"title": p["title"], "url": p["url"], "nick": p["nick_name"], "cover": None}
        if p["cover_url"]:
            rec["cover"] = _dl(p["cover_url"], "xiaohongshu", base / "cover.jpg", pub + "/cover.jpg")
        recs.append(rec)
    return {"ok": recs, "error": None}


def _prev_douyin() -> dict:
    """抖音关闭时，从上次的 search_results.json 沿用抖音结果（按 query 文本取），保住已抓到的视频。"""
    if not OUT.exists():
        return {}
    try:
        prev = json.loads(OUT.read_text("utf-8"))
    except Exception:
        return {}
    return {r["query"]: r.get("douyin") for r in prev if r.get("query")}


def main() -> None:
    settings = Settings.from_env()
    demo = json.loads(DEMO.read_text("utf-8"))
    # 三平台搜索各自独立闸（随机 10~12s）；抖音详情更轻（3~5s）
    dy_search = RateLimiter(min_interval_seconds=SEARCH_MIN, max_interval_seconds=SEARCH_MAX)
    dy_detail = RateLimiter(min_interval_seconds=3.0, max_interval_seconds=5.0)
    wx_search = RateLimiter(min_interval_seconds=SEARCH_MIN, max_interval_seconds=SEARCH_MAX)
    xhs_search = RateLimiter(min_interval_seconds=SEARCH_MIN, max_interval_seconds=SEARCH_MAX)
    prev_dy = {} if DOUYIN_ON else _prev_douyin()
    if not DOUYIN_ON:
        print(f"（抖音已关闭：沿用上次 {len(prev_dy)} 条 query 的抖音结果，本次只跑微信+小红书）")
    results = []
    for tac in ("tactic1", "tactic2", "tactic4"):
        t = demo.get(tac) or {}
        name = t.get("name", tac)
        for it in (t.get("items") or [])[:N]:
            q = it.get("query")
            if not q:
                continue
            print(f"[{name}] {q}")
            if DOUYIN_ON:
                dy = douyin_topk(q, settings, dy_search, dy_detail)
            else:
                dy = prev_dy.get(q) or {"ok": [], "error": "抖音限流暂停，未跑"}
            wx = weixin_topk(q, settings, wx_search)
            xhs = xiaohongshu_topk(q, settings, xhs_search)
            print(f"   抖音 {len(dy['ok'])} 视频{_err(dy)}"
                  f"   微信 {len(wx['ok'])} 封面{_err(wx)}"
                  f"   小红书 {len(xhs['ok'])} 封面{_err(xhs)}")
            results.append({"method": name, "query": q, "douyin": dy, "weixin": wx, "xiaohongshu": xhs})
    OUT.parent.mkdir(parents=True, exist_ok=True)
    OUT.write_text(json.dumps(results, ensure_ascii=False, indent=1), encoding="utf-8")
    tot = lambda k: sum(len(r[k]["ok"]) for r in results)
    print(f"\nwrote {len(results)} 条 query（抖音 {tot('douyin')} / 微信 {tot('weixin')} / "
          f"小红书 {tot('xiaohongshu')}）→ {OUT}")


def _err(d: dict) -> str:
    return " / " + d["error"] if d.get("error") else ""


if __name__ == "__main__":
    main()