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

"""帖子级「创作知识 / 非创作知识」分类 + 提取知识点。

判据是「内容作品创作知识」，覆盖图文、视频、脚本、游戏视频、历史视频等内容创作场景，
并显式排除三类越界（见 prompts）：① 应试/学术写作
② 制作/工具操作（=制作知识）③ 学科知识/评论/作品本身。两套真实提示词：
  图文（小红书/微信）：prompts/classify_imgtext.txt（标题+正文+全部图喂多模态模型）
  视频（抖音）：prompts/classify_video.txt（OSS/CDN URL 直喂；本地大视频 fallback 时先压 480p 保音轨）
都输出 {is_empty, reason, knowledge}：is_empty 即创作闸；is_empty=false 时连带提炼出具体创作知识点。
结果按 url 写 app.db 的 post_class（upsert 覆盖，可重判）。并发；与微信补下载并行安全（busy_timeout）。
只读 prompts、走 OpenRouter / Ark / Qwen——不改 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 os
import subprocess
import sys
import tempfile
import threading
import time
from pathlib import Path
from urllib.parse import urlparse

import httpx

from acquisition import store
from core.config import Settings, load_env_file
from core.prompts import load_prompt

ROOT = Path(__file__).resolve().parent.parent
PLATFORMS = ["xiaohongshu", "weixin", "douyin"]   # 默认重判全部；可传平台名覆盖
IMG_WORKERS = 8                          # 图文并发
VID_WORKERS = 3                          # 视频并发（含 ffmpeg 压制，别太高）
MAX_CARDS = 12                           # 图文最多送几张图
COMPRESS_OVER_MB = 12                    # 视频超过此大小先压再喂
COMPRESS_H = 480
DEFAULT_PROVIDER_MIN_INTERVAL_SECONDS = 0.5
DEFAULT_429_BACKOFF_SECONDS = [30.0, 90.0, 180.0]
RETRYABLE_STATUS_CODES = {429, 500, 502, 503, 504}


class _ProviderThrottle:
    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._next_allowed = 0.0

    def wait(self, min_interval_seconds: float) -> None:
        with self._lock:
            now = time.monotonic()
            sleep_for = max(0.0, self._next_allowed - now)
            self._next_allowed = max(now, self._next_allowed) + min_interval_seconds
        if sleep_for > 0:
            time.sleep(sleep_for)

    def penalize(self, seconds: float) -> None:
        if seconds <= 0:
            return
        with self._lock:
            self._next_allowed = max(self._next_allowed, time.monotonic() + seconds)


_PROVIDER_THROTTLES: dict[str, _ProviderThrottle] = {}
_PROVIDER_THROTTLES_LOCK = threading.Lock()


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 _is_http_url(value: str) -> bool:
    try:
        parsed = urlparse(value)
    except Exception:
        return False
    return parsed.scheme in ("http", "https") and bool(parsed.netloc)


def _compress(mp4: Path) -> Path:
    """大视频压到 480p（保留口播音轨）→ 临时 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 _parse_judge_content(content: str) -> tuple:
    d = json.loads(content)
    if bool(d.get("is_empty")):
        return 0, str(d.get("reason", ""))[:60], "", ""
    return 1, str(d.get("reason", ""))[:60], str(d.get("knowledge", "") or ""), ""


def _env_first(env: dict, *keys: str) -> str:
    for key in keys:
        value = os.getenv(key) or env.get(key)
        if value:
            return value
    return ""


def _is_qwen_model(model: str) -> bool:
    lower = model.lower()
    return lower.startswith("qwen") or lower.startswith("qwq")


def _provider_key(name: str) -> str:
    return name.split(":", 1)[0].lower()


def _provider_throttle(provider: str) -> _ProviderThrottle:
    with _PROVIDER_THROTTLES_LOCK:
        throttle = _PROVIDER_THROTTLES.get(provider)
        if throttle is None:
            throttle = _ProviderThrottle()
            _PROVIDER_THROTTLES[provider] = throttle
        return throttle


def _env_float(env: dict, default: float, *keys: str) -> float:
    value = _env_first(env, *keys)
    if not value:
        return default
    try:
        return max(0.0, float(value))
    except ValueError:
        return default


def _provider_min_interval(provider: str, env: dict) -> float:
    prefix = provider.upper()
    return _env_float(
        env,
        DEFAULT_PROVIDER_MIN_INTERVAL_SECONDS,
        f"CLASSIFY_{prefix}_MIN_INTERVAL_SECONDS",
        "CLASSIFY_PROVIDER_MIN_INTERVAL_SECONDS",
    )


def _parse_backoff_list(value: str) -> list[float]:
    out = []
    for part in value.split(","):
        try:
            seconds = float(part.strip())
        except ValueError:
            continue
        if seconds > 0:
            out.append(seconds)
    return out or DEFAULT_429_BACKOFF_SECONDS


def _retry_after_seconds(resp: httpx.Response) -> float | None:
    value = resp.headers.get("retry-after")
    if not value:
        return None
    try:
        return max(0.0, float(value))
    except ValueError:
        return None


def _provider_429_backoff(provider: str, env: dict, attempt: int, resp: httpx.Response) -> float:
    retry_after = _retry_after_seconds(resp)
    if retry_after is not None:
        return retry_after
    prefix = provider.upper()
    values = _parse_backoff_list(
        _env_first(
            env,
            f"CLASSIFY_{prefix}_429_BACKOFF_SECONDS",
            "CLASSIFY_429_BACKOFF_SECONDS",
        ) or ",".join(str(v) for v in DEFAULT_429_BACKOFF_SECONDS)
    )
    return values[min(attempt, len(values) - 1)]


def _providers(settings: Settings, messages: list) -> list[tuple[str, str, dict, dict]]:
    body = {"model": settings.video_model, "messages": messages,
            "response_format": {"type": "json_object"}}
    env = load_env_file(os.getenv("CK_ENV_FILE", ".env"))
    out: list[tuple[str, str, dict, dict]] = []
    provider = (_env_first(env, "CLASSIFY_PROVIDER") or "auto").lower()
    explicit_model = _env_first(env, "CLASSIFY_MODEL")

    bailian_key = _env_first(
        env,
        "ALIYUN_BAILIAN_API_KEY",
    )
    if bailian_key and provider in ("auto", "qwen", "bailian"):
        bailian_url = _env_first(
            env,
            "ALIYUN_BAILIAN_BASE_URL",
        ) or "https://dashscope.aliyuncs.com/compatible-mode/v1"
        models = []
        for model in [
            explicit_model,
            _env_first(env, "ALIYUN_BAILIAN_MODEL"),
            _env_first(env, "OPENROUTER_MODEL"),
            "qwen3.7-plus",
            "qwen-vl-plus",
        ]:
            if model and _is_qwen_model(model) and model not in models:
                models.append(model)
        for model in models:
            out.append((
                f"qwen:{model}",
                bailian_url.rstrip("/") + "/chat/completions",
                {"Authorization": f"Bearer {bailian_key}", "Content-Type": "application/json"},
                {"model": model, "messages": messages, "response_format": {"type": "json_object"}},
            ))
        if provider in ("qwen", "bailian"):
            return out

    if settings.openrouter_api_key and provider in ("auto", "openrouter"):
        out.append((
            "openrouter",
            settings.openrouter_base_url.rstrip("/") + "/chat/completions",
            {"Authorization": f"Bearer {settings.openrouter_api_key}", "Content-Type": "application/json"},
            body,
        ))

    ark_key = _env_first(env, "ARK_API_KEY")
    if ark_key and provider in ("auto", "ark"):
        ark_url = _env_first(env, "ARK_CHAT_URL") or "https://ark.cn-beijing.volces.com/api/v3/chat/completions"
        models = []
        explicit = explicit_model or _env_first(env, "ARK_CHAT_MODEL")
        if explicit:
            models.append(explicit)
        # Seed 2 Mini 接入点/1.6 Vision 更适合图文+视频理解；flash 作为轻量兜底。
        models.extend(["ep-20260506151915-jqvw7", "doubao-seed-1-6-vision-250815", "doubao-seed-1-6-flash-250615"])
        seen = set()
        for model in models:
            if model in seen:
                continue
            seen.add(model)
            out.append((
                f"ark:{model}",
                ark_url,
                {"Authorization": f"Bearer {ark_key}", "Content-Type": "application/json"},
                {"model": model, "messages": messages, "response_format": {"type": "json_object"}},
            ))
    return out


def _judge(messages: list, settings: Settings, timeout: float) -> tuple:
    """调多模态模型（Qwen / OpenRouter / Ark），解析 {is_empty, reason, knowledge}，带重试。
    返回 (is_creation 1/0/None, reason, knowledge, points)。"""
    last = ""
    env = load_env_file(os.getenv("CK_ENV_FILE", ".env"))
    for name, api, headers, payload in _providers(settings, messages):
        provider = _provider_key(name)
        throttle = _provider_throttle(provider)
        min_interval = _provider_min_interval(provider, env)
        for attempt in range(3):
            try:
                throttle.wait(min_interval)
                resp = httpx.post(api, headers=headers, json=payload, timeout=timeout)
                if resp.status_code == 200:
                    return _parse_judge_content(resp.json()["choices"][0]["message"]["content"])
                last = f"{name} http {resp.status_code}"
                if resp.status_code == 429:
                    throttle.penalize(_provider_429_backoff(provider, env, attempt, resp))
                if resp.status_code in (401, 403, 404):
                    break
            except Exception as exc:
                last = f"{name} {str(exc)[:50]}"
            if "http " not in last or any(f"http {code}" in last for code in RETRYABLE_STATUS_CODES):
                time.sleep(2 * (attempt + 1))
            else:
                break
    if not last:
        last = "missing Qwen/OpenRouter/Ark credentials"
    return None, f"判定失败: {last}", "", ""


def classify_imgtext(p: dict, settings: Settings) -> tuple:
    """图文：标题+正文+全部图，用收紧的 classify_imgtext.txt 判 is_empty 并提取知识点。"""
    user = [{"type": "text", "text": f"平台：{p.get('platform')}\n标题：{p.get('title', '')}\n"
             f"正文：{(p.get('body_text') or '')[:1500]}\n（下附帖子图片，请一并看完）"}]
    for im in (p.get("images") or [])[:MAX_CARDS]:
        u = _data_url(im)
        if u:
            user.append({"type": "image_url", "image_url": {"url": u}})
    messages = [{"role": "system", "content": load_prompt("classify_imgtext")},
                {"role": "user", "content": user}]
    return _judge(messages, settings, timeout=120)


def classify_video(p: dict, settings: Settings) -> tuple:
    """视频：看完整段视频，用 classify_video.txt 判 is_empty 并提取知识点。

    新采集链路里抖音视频会先转存 OSS，传入 HTTP(S) CDN URL；老链路传 /data 本地 mp4。
    """
    rel = p.get("video") or ""
    if _is_http_url(rel):
        media = rel
    else:
        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}}]}]
    return _judge(messages, settings, timeout=300)


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()
    conn = store.connect()
    platforms = sys.argv[1:] or PLATFORMS
    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, settings): 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']}（本轮失败 {done['fail']}，可重跑补判）")


if __name__ == "__main__":
    main()