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

"""把已下载的帖子分成「创作知识 / 非创作知识」（只判定，不解构、不提知识点）。

复用 prompts/extract.txt 里的「创作知识」定义（只读，绝不改 skill/creation_knowledge）。
方案 2：标题 + 正文 + 封面/帖内图一起喂 Gemini（OpenRouter 多模态），并发 10 + 失败重试。
结果写 data/app.db 的 post_class 表（按 url），前端按 url 显示「创作知识 / 非创作知识」角标。
不碰视频/帖子文件本身，不入 ingest。重跑只补未分类的（幂等）。
用法：PYTHONPATH=. CK_ENV_FILE=.env python -m acquisition.classify
"""
from __future__ import annotations

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

import httpx

from acquisition import store
from core.config import Settings

ROOT = Path(__file__).resolve().parent.parent
EXTRACT_PROMPT = ROOT / "prompts" / "extract.txt"
WORKERS = 10          # 并发路数（实测线性提速；过高易被 OpenRouter 限流）
MAX_IMG = 2           # 每帖最多喂几张图（控延迟/体积）
RETRIES = 3           # 失败重试次数（含 429 限流退避）


def _definition() -> str:
    """复用 extract.txt 的 <什么算创作知识> 段（同一判据，只读不改）。"""
    txt = EXTRACT_PROMPT.read_text("utf-8")
    a, b = txt.find("<什么算创作知识>"), txt.find("</什么算创作知识>")
    return txt[a:b + len("</什么算创作知识>")] if a >= 0 and b >= 0 else ""


SYS_TMPL = """你是创作知识分类器。给你一篇帖子的平台/标题/正文和图片，判断它是不是「创作知识」。
{definition}
只输出 JSON：{{"is_creation": true/false, "reason": "一句话理由（≤25字）"}}。
范围内但拿不准 → is_creation=true；明确越界（游戏玩法/产品运营/做菜本身/纯工具操作/纯作品展示）→ false。
注意：视频帖可能只有标题+封面、信息有限，据现有内容审慎判断。"""


def _img_data_url(public_path: str):
    """本地公开路径 /data/search/... → base64 data URL；不存在返回 None。"""
    fs = ROOT / public_path.lstrip("/")
    if not fs.exists():
        return None
    return "data:image/jpeg;base64," + base64.b64encode(fs.read_bytes()).decode()


def classify_one(post: dict, settings: Settings, system: str) -> dict:
    """单帖多模态判定，带重试。返回 {is_creation:1/0/None, reason}。None=判定失败（留待重跑）。"""
    user_text = (f"平台：{post.get('platform')}\n标题：{post.get('title') or ''}\n"
                 f"正文：{(post.get('body_text') or '')[:600]}")
    content = [{"type": "text", "text": user_text}]
    for img in (post.get("images") or [])[:MAX_IMG]:
        url = _img_data_url(img)
        if url:
            content.append({"type": "image_url", "image_url": {"url": url}})
    messages = [{"role": "system", "content": system}, {"role": "user", "content": content}]
    api = settings.openrouter_base_url.rstrip("/") + "/chat/completions"
    headers = {"Authorization": f"Bearer {settings.openrouter_api_key}", "Content-Type": "application/json"}
    payload = {"model": settings.llm_model, "messages": messages,
               "response_format": {"type": "json_object"}}
    last = ""
    for attempt in range(RETRIES):
        try:
            resp = httpx.post(api, headers=headers, json=payload, timeout=90)
            if resp.status_code == 200:
                d = json.loads(resp.json()["choices"][0]["message"]["content"])
                return {"is_creation": 1 if d.get("is_creation") else 0,
                        "reason": str(d.get("reason", ""))[:50]}
            last = f"http {resp.status_code}"
        except Exception as exc:
            last = str(exc)[:50]
        time.sleep(1.5 * (attempt + 1))     # 退避后重试
    return {"is_creation": None, "reason": f"判定失败: {last}"}


def main() -> None:
    settings = Settings.from_env()
    system = SYS_TMPL.format(definition=_definition())
    conn = store.connect()
    posts = store.get_unclassified_posts(conn)
    if not posts:
        print("没有待分类帖子（都判过了）。"); conn.close(); return
    print(f"待分类 {len(posts)} 帖（并发 {WORKERS}，每帖≤{MAX_IMG} 图）...")
    ts, done, fail = int(time.time()), 0, 0
    with cf.ThreadPoolExecutor(max_workers=WORKERS) as ex:
        futs = {ex.submit(classify_one, p, settings, system): p for p in posts}
        for fut in cf.as_completed(futs):
            p = futs[fut]
            r = fut.result()
            if r["is_creation"] is None:
                fail += 1
            else:
                store.upsert_class(conn, p["url"], r["is_creation"], r["reason"], ts)
            done += 1
            if done % 25 == 0:
                print(f"  {done}/{len(posts)}（失败 {fail}）")
    c = store.class_counts(conn)
    conn.close()
    print(f"完成：创作知识 {c['creation']} / 非创作知识 {c['non_creation']}（本轮失败 {fail}，可重跑补判）")


if __name__ == "__main__":
    main()