zhiwei-ima-memory/scripts/auto_classify.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
auto_classify.py —— 知识库公众号文章自动归类
============================================
把知识库根目录下「未归类的微信公众号文章」按标题关键词匹配，
自动移动到对应的主题目录（文件夹）。

用法：
    python3 auto_classify.py            # 运行一次（实际移动）
    python3 auto_classify.py --dry-run  # 仅预览匹配结果，不移动
    python3 auto_classify.py --loop 10  # 每 10 分钟循环扫描一次
    python3 auto_classify.py --once     # 等价默认（运行一次）

依赖环境变量：
    IMA_OPENAPI_CLIENTID
    IMA_OPENAPI_APIKEY
知识库 ID 见下方 KB_ID（源的个人知识库，公众号文章自动同步到此）。
"""
import os
import sys
import json
import time
import argparse
import urllib.request
from collections import defaultdict

BASE_URL = "https://ima.qq.com"
KB_ID = "fsUjWTJ9mncR91ppaaPlxW2ibb0scLyM80jHHvyN0jI="
LOG_PATH = "/sandbox/workspace/auto_classify.log"
BATCH = 10  # 单次 move_knowledge 最多 10 个

# 关键词 → 目标文件夹（按列表顺序优先级匹配，先命中先得）。
# folder 名称需与知识库内真实文件夹名一致；不存在的会回落到 DEFAULT_FOLDER。
RULES = [
    ("HermesAgent",        ["Hermes", "爱马仕", "hermes"]),
    ("ClaudeCode",         ["Claude", "Anthropic", "Codex", "claude"]),
    ("OpenClaw",           ["OpenClaw", "龙虾", "Claw", "openclaw"]),
    ("DeepSeekV4",         ["DeepSeek", "Reasonix", "deepseek"]),
    ("AIAgent智能体",       ["Agent", "智能体", "OpenHuman", "agent"]),
    ("MCP模型上下文协议",    ["MCP", "AI网关", "mcp"]),
    ("阿里AI生态",          ["阿里", "Qwen", "通义", "qwen", "千问"]),
    ("LLM大模型",           ["LLM", "大模型", "年入", "程序员", "创始人", "llm"]),
    ("CoWork",             ["Workbuddy", "CoWork", "workbuddy"]),
    ("Skill本地技能",        ["Skill", "技能", "skill", "仓颉"]),
    ("AI自然语言编程",       ["编程", "前端", "后端", "TypeScript", "React"]),
    ("github开源工具",      ["github", "GitHub", "开源", "标星", "star"]),
    ("智普GLM",            ["智普", "GLM", "glm", "ChatGLM"]),
    ("智能化运维AIOPS",      ["Token工厂", "Token成本", "算力", "运维", "AIOps", "aiops", "可观测"]),
    ("Office",             ["Office", "出片", "剪辑", "视频生成", "Excel", "PPT", "Word", "Obsidian"]),
    ("RAG知识库应用",        ["RAG", "rag", "知识库", "Karpathy", "检索"]),
    ("AI大会和论坛",        ["大会", "论坛", "峰会", "全球脑力"]),
    ("语音ASR小智TTS",      ["语音", "ASR", "TTS", "有声书", "电子书"]),
    ("多模态",              ["多模态", "VL2", "SAIL", "视觉"]),
    ("机器人",              ["机器人", "ASAP", "仿真"]),
    ("音乐AI",             ["音乐", "SongBloom", "歌曲生成"]),
    ("翻译工具",            ["翻译", "TinyAI"]),
]

# 命中以下标记的文章不参与归类（微信平台类、致谢类等）
SKIP_MARKERS = ["微信公众平台", "对联", "谢谢", "__SKIP__"]
DEFAULT_FOLDER = "其他"


def log(msg):
    ts = time.strftime("%Y-%m-%d %H:%M:%S")
    line = f"[{ts}] {msg}"
    print(line)
    try:
        with open(LOG_PATH, "a", encoding="utf-8") as f:
            f.write(line + "\n")
    except Exception:
        pass


def call_api(endpoint, body):
    url = BASE_URL + endpoint
    data = json.dumps(body).encode("utf-8")
    req = urllib.request.Request(url, data=data, method="POST")
    req.add_header("ima-openapi-clientid", os.environ.get("IMA_OPENAPI_CLIENTID", ""))
    req.add_header("ima-openapi-apikey", os.environ.get("IMA_OPENAPI_APIKEY", ""))
    req.add_header("Content-Type", "application/json")
    with urllib.request.urlopen(req, timeout=30) as resp:
        r = json.loads(resp.read().decode("utf-8"))
        return r.get("data", r)


def list_all(kb_id, folder_id=None):
    items = []
    cursor = ""
    while True:
        body = {"cursor": cursor, "limit": 50, "knowledge_base_id": kb_id}
        if folder_id:
            body["folder_id"] = folder_id
        r = call_api("/openapi/wiki/v1/get_knowledge_list", body)
        items.extend(r.get("knowledge_list", []))
        if r.get("is_end") or not r.get("next_cursor"):
            break
        cursor = r["next_cursor"]
    return items


def get_folder_map(kb_id):
    items = list_all(kb_id)
    return {it["title"]: it["media_id"] for it in items if it.get("media_type") == 99}


def classify(title, folder_map):
    low = title.lower()
    for folder, kws in RULES:
        for kw in kws:
            if kw.lower() in low:
                if folder in folder_map:
                    return folder, folder_map[folder]
    if DEFAULT_FOLDER in folder_map:
        return DEFAULT_FOLDER, folder_map[DEFAULT_FOLDER]
    return None, None


def run_once(dry_run):
    folder_map = get_folder_map(KB_ID)
    items = list_all(KB_ID)
    articles = [it for it in items if it.get("media_type") == 6]  # 公众号文章
    log(f"根目录共 {len(items)} 项，公众号文章 {len(articles)} 篇；文件夹 {len(folder_map)} 个")

    moves = []  # (article, target_name, target_id)
    skipped = 0
    for a in articles:
        title = a.get("title", "")
        if any(m in title for m in SKIP_MARKERS):
            skipped += 1
            continue
        name, fid = classify(title, folder_map)
        if name is None:
            continue  # 无兜底文件夹，留在根目录
        moves.append((a, name, fid))

    if dry_run:
        for a, name, _ in moves:
            log(f"[DRY-RUN] «{a['title']}» -> 「{name}」")
        log(f"[DRY-RUN] 预览：将移动 {len(moves)} 篇，跳过 {skipped} 篇")
        return len(moves)

    groups = defaultdict(list)
    for a, name, fid in moves:
        groups[fid].append((a["media_id"], name, a["title"]))
    total = 0
    for fid, batch in groups.items():
        for i in range(0, len(batch), BATCH):
            chunk = batch[i:i + BATCH]
            infos = [{"media_id": m} for m, _, _ in chunk]
            try:
                call_api("/openapi/wiki/v1/move_knowledge", {
                    "src_knowledge_base_id": KB_ID,
                    "dst_knowledge_base_id": KB_ID,
                    "dst_folder_id": fid,
                    "infos": infos,
                })
                total += len(chunk)
                for _, name, title in chunk:
                    log(f"✓ «{title}» -> 「{name}」")
            except Exception as e:
                log(f"✗ 移动失败: {e}")
    log(f"本次移动 {total} 篇，跳过 {skipped} 篇")
    return total


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--dry-run", action="store_true", help="仅预览匹配结果，不移动")
    ap.add_argument("--loop", type=int, default=0, help="循环间隔分钟，0=不循环")
    ap.add_argument("--once", action="store_true", help="只运行一次（默认）")
    args = ap.parse_args()

    if not os.environ.get("IMA_OPENAPI_CLIENTID") or not os.environ.get("IMA_OPENAPI_APIKEY"):
        log("✗ 缺少 IMA_OPENAPI_CLIENTID / IMA_OPENAPI_APIKEY 环境变量")
        sys.exit(1)

    if args.loop:
        interval = args.loop * 60
        log(f"循环模式启动，每 {args.loop} 分钟扫描一次")
        while True:
            try:
                run_once(args.dry_run)
            except Exception as e:
                log(f"✗ 运行异常: {e}")
            time.sleep(interval)
    else:
        run_once(args.dry_run)


if __name__ == "__main__":
    main()