Create-knowledge-find-decode/acquisition/classify.py

"""帖子级「创作知识 / 非创作知识」分类——忠实复用 pipeline 的「读懂 + 创作闸」，不简化提示词。

图文（小红书/微信）：用 creation_knowledge 的 GeminiExtractor（**完整 prompts/extract.txt**，
  带【卡片N】把全部图喂给 Gemini）→ ExtractedContent.is_empty 即创作闸（true=非创作）。
视频（抖音）：用 extract_video（**完整 prompts/extract_video.txt**，整段视频原生喂 Gemini）→
  is_empty（=看完视频没提到任何创作知识段落）。大视频先 ffmpeg 压到 480p（保留口播音轨）再 base64。
结果按 url 写 app.db 的 post_class（upsert 覆盖，可重判）。并发；与微信补下载并行安全（busy_timeout）。
只读 prompts、import 调用 pipeline 函数——不改 skill/creation_knowledge，不解构、不入 ingest。
用法：PYTHONPATH=. CK_ENV_FILE=.env python -m acquisition.classify
"""
from __future__ import annotations

import base64
import concurrent.futures as cf
import json
import subprocess
import sys
import tempfile
import time
from pathlib import Path

import httpx

from acquisition import store
from core.config import Settings
from core.models import Card, Post
from core.prompts import load_prompt
from creation_knowledge.integrations.extractor import GeminiExtractor

ROOT = Path(__file__).resolve().parent.parent
PLATFORMS = ["xiaohongshu", "douyin"]    # 本轮重判平台（微信等正文下完再单独跑）
IMG_WORKERS = 8                          # 图文并发
VID_WORKERS = 3                          # 视频并发（含 ffmpeg 压制，别太高）
MAX_CARDS = 12                           # 图文最多送几张图（与 extractor 对齐）
COMPRESS_OVER_MB = 12                    # 视频超过此大小先压再喂
COMPRESS_H = 480


def _data_url(public_path: str):
    fs = ROOT / public_path.lstrip("/")
    if not fs.exists():
        return None
    return "data:image/jpeg;base64," + base64.b64encode(fs.read_bytes()).decode()


def _imgtext_post(p: dict) -> Post:
    cards = []
    for i, im in enumerate((p.get("images") or [])[:MAX_CARDS], start=1):
        u = _data_url(im)
        if u:
            cards.append(Card(index=i, kind="image", url=u))
    return Post(id=f"{p['platform']}_cls", platform=p["platform"], url=p["url"],
                content_id=p["url"], title=p.get("title", ""),
                body_text=p.get("body_text", ""), cards=cards)


def _compress(mp4: Path) -> Path:
    """大视频压到 480p（保留音轨——extract_video 要听口播）→ 临时 mp4；失败回原文件。"""
    out = Path(tempfile.gettempdir()) / f"ck_{mp4.parent.name}_480.mp4"
    try:
        import imageio_ffmpeg
        ff = imageio_ffmpeg.get_ffmpeg_exe()
        subprocess.run([ff, "-y", "-i", str(mp4), "-vf", f"scale=-2:{COMPRESS_H}",
                        "-c:v", "libx264", "-crf", "30", "-preset", "veryfast",
                        "-c:a", "aac", "-b:a", "64k", str(out)],
                       capture_output=True, timeout=300)
        if out.exists() and out.stat().st_size > 0:
            return out
    except Exception:
        pass
    return mp4


def classify_imgtext(p: dict, extractor: GeminiExtractor) -> tuple:
    ec = extractor.extract(_imgtext_post(p))           # 完整 extract.txt → is_empty + 提取的知识
    if ec.is_empty:
        return 0, "判为非创作：无可迁移的创作方法", "", ""
    points = json.dumps([{"index": c.index, "content": c.content} for c in (ec.cards or [])],
                        ensure_ascii=False)
    return 1, (ec.text or "")[:60], (ec.text or ""), points   # reason 取开头，knowledge 存全文


def classify_video(p: dict, settings: Settings) -> tuple:
    """看完整段视频，用与图文同款的【严格创作知识判据】（classify_video.txt：两轴+越界+拦观点拔高）判 is_empty。
    不再用 extract_video 的松「提炼」闸（它会把讲观点的视频拔高成创作知识）。"""
    rel = p.get("video") or ""
    mp4 = ROOT / rel.lstrip("/")
    if not rel or not mp4.exists():
        return None, "无本地视频", "", ""
    use = _compress(mp4) if mp4.stat().st_size > COMPRESS_OVER_MB * 1048576 else mp4
    try:
        media = "data:video/mp4;base64," + base64.b64encode(use.read_bytes()).decode()
    finally:
        if use != mp4:
            try:
                use.unlink()
            except Exception:
                pass
    messages = [{"role": "system", "content": load_prompt("classify_video")},
                {"role": "user", "content": [{"type": "text", "text": "判断这条视频是不是创作知识。"},
                                             {"type": "video_url", "video_url": {"url": media}}]}]
    api = settings.openrouter_base_url.rstrip("/") + "/chat/completions"
    headers = {"Authorization": f"Bearer {settings.openrouter_api_key}", "Content-Type": "application/json"}
    payload = {"model": settings.video_model, "messages": messages,
               "response_format": {"type": "json_object"}}
    last = ""
    for attempt in range(3):
        try:
            resp = httpx.post(api, headers=headers, json=payload, timeout=300)
            if resp.status_code == 200:
                d = json.loads(resp.json()["choices"][0]["message"]["content"])
                is_empty = bool(d.get("is_empty"))
                if is_empty:
                    return 0, str(d.get("reason", ""))[:60], "", ""
                kn = str(d.get("knowledge", "") or "")
                return 1, str(d.get("reason", ""))[:60], kn, ""
            last = f"http {resp.status_code}"
        except Exception as exc:
            last = str(exc)[:50]
        time.sleep(2 * (attempt + 1))
    return None, f"判定失败: {last}", "", ""


def _safe(fn, *a) -> tuple:
    try:
        return fn(*a)
    except Exception as exc:
        return None, f"判定失败: {str(exc)[:60]}", "", ""


def main() -> None:
    settings = Settings.from_env()
    extractor = GeminiExtractor.from_env()
    conn = store.connect()
    platforms = sys.argv[1:] or PLATFORMS    # 可传平台名重判，如：... -m acquisition.classify weixin
    posts = store.posts_to_classify(conn, platforms)
    imgs = [p for p in posts if p["platform"] != "douyin"]
    vids = [p for p in posts if p["platform"] == "douyin"]
    total = len(posts)
    print(f"忠实重判：图文 {len(imgs)}（并发{IMG_WORKERS}）+ 抖音视频 {len(vids)}（并发{VID_WORKERS}，大视频先压）")
    ts = int(time.time())
    done = {"n": 0, "fail": 0}

    def _write(p, res):
        ic, reason, knowledge, points = res
        if ic is None:
            done["fail"] += 1
        else:
            store.upsert_class(conn, p["url"], ic, reason, ts, knowledge, points)
        done["n"] += 1
        if done["n"] % 20 == 0:
            print(f"  {done['n']}/{total}（失败 {done['fail']}）")

    with cf.ThreadPoolExecutor(IMG_WORKERS) as ex:
        futs = {ex.submit(_safe, classify_imgtext, p, extractor): p for p in imgs}
        for fut in cf.as_completed(futs):
            _write(futs[fut], fut.result())
    with cf.ThreadPoolExecutor(VID_WORKERS) as ex:
        futs = {ex.submit(_safe, classify_video, p, settings): p for p in vids}
        for fut in cf.as_completed(futs):
            _write(futs[fut], fut.result())

    c = store.class_counts(conn)
    conn.close()
    print(f"完成（含历史微信）：创作知识 {c['creation']} / 非创作知识 {c['non_creation']}"
          f"（本轮失败 {done['fail']}，可重跑补判）")


if __name__ == "__main__":
    main()