Create-knowledge-find-decode/acquisition/classification/coarse.py

"""Formal coarse classifier for creation-knowledge candidates."""
from __future__ import annotations

import hashlib
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from acquisition.classify import (
    MAX_CARDS,
    _data_url,
    _is_http_url,
    _judge,
    classify_video as _legacy_classify_video,
)
from core.config import Settings
from core.prompts import load_prompt

ROOT = Path(__file__).resolve().parents[2]


@dataclass(frozen=True)
class ClassificationResult:
    is_creation_knowledge: bool | None
    label: str | None
    confidence: float | None
    reason: str
    knowledge: str = ""
    prompt_version: str | None = None
    result_payload: dict[str, Any] = field(default_factory=dict)
    status: str = "classified"
    error_message: str | None = None


def prompt_version(*names: str) -> str:
    h = hashlib.sha256()
    for name in names:
        path = ROOT / "prompts" / f"{name}.txt"
        h.update(name.encode("utf-8"))
        if path.exists():
            h.update(path.read_bytes())
    return h.hexdigest()[:16]


def classify_imgtext(payload: dict[str, Any], settings: Settings) -> tuple:
    """Classify image-text content, accepting both HTTP image URLs and /data paths."""
    user = [
        {
            "type": "text",
            "text": (
                f"平台：{payload.get('platform')}\n"
                f"标题：{payload.get('title', '')}\n"
                f"正文：{(payload.get('body_text') or '')[:1500]}\n"
                "（下附帖子图片，请一并看完）"
            ),
        }
    ]
    for image in (payload.get("images") or [])[:MAX_CARDS]:
        if _is_http_url(image):
            user.append({"type": "image_url", "image_url": {"url": image}})
            continue
        data_url = _data_url(image, settings)
        if data_url:
            user.append({"type": "image_url", "image_url": {"url": data_url}})
    messages = [
        {"role": "system", "content": load_prompt("classify_imgtext")},
        {"role": "user", "content": user},
    ]
    return _judge(messages, settings, timeout=120)


def classify_video(payload: dict[str, Any], settings: Settings) -> tuple:
    return _legacy_classify_video(payload, settings)


def coarse_classify_item(
    *,
    platform: str,
    title: str = "",
    body_text: str = "",
    image_urls: list[str] | None = None,
    video_url: str = "",
    settings: Settings,
) -> ClassificationResult:
    if platform == "douyin" or video_url:
        version = prompt_version("classify_video")
        is_creation, reason, knowledge, points = classify_video(
            {
                "platform": platform,
                "title": title,
                "body_text": body_text,
                "video": video_url,
            },
            settings,
        )
    else:
        version = prompt_version("classify_imgtext")
        is_creation, reason, knowledge, points = classify_imgtext(
            {
                "platform": platform,
                "title": title,
                "body_text": body_text,
                "images": image_urls or [],
            },
            settings,
        )

    if is_creation is None:
        return ClassificationResult(
            is_creation_knowledge=None,
            label=None,
            confidence=None,
            reason=reason,
            knowledge=knowledge,
            prompt_version=version,
            result_payload={"knowledge": knowledge, "points": points},
            status="failed",
            error_message=reason,
        )
    is_hit = bool(is_creation)
    return ClassificationResult(
        is_creation_knowledge=is_hit,
        label="creation" if is_hit else "not_creation",
        confidence=1.0,
        reason=reason,
        knowledge=knowledge,
        prompt_version=version,
        result_payload={"knowledge": knowledge, "points": points},
        status="classified",
    )