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
from core.text_limits import CLASSIFY_BODY_MAX_CHARS, clip_text
from pipeline.tracing import TraceContext, TraceWriter

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,
    *,
    trace_writer: TraceWriter | None = None,
    trace_context: TraceContext | None = None,
) -> 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"正文：{clip_text(payload.get('body_text') or '', CLASSIFY_BODY_MAX_CHARS)}\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},
    ]
    if trace_writer is None and trace_context is None:
        return _judge(messages, settings, timeout=120)
    return _judge(
        messages,
        settings,
        timeout=120,
        trace_writer=trace_writer,
        trace_context=trace_context,
        trace_stage="classify",
        trace_substage="coarse_imgtext",
        prompt_name="classify_imgtext",
    )


def classify_video(
    payload: dict[str, Any],
    settings: Settings,
    *,
    trace_writer: TraceWriter | None = None,
    trace_context: TraceContext | None = None,
) -> tuple:
    if trace_writer is None and trace_context is None:
        return _legacy_classify_video(payload, settings)
    return _legacy_classify_video(
        payload,
        settings,
        trace_writer=trace_writer,
        trace_context=trace_context,
    )


def coarse_classify_item(
    *,
    platform: str,
    content_mode: str | None = None,
    title: str = "",
    body_text: str = "",
    image_urls: list[str] | None = None,
    video_url: str = "",
    settings: Settings,
    trace_writer: TraceWriter | None = None,
    trace_context: TraceContext | None = None,
) -> ClassificationResult:
    if content_mode == "unsupported":
        return ClassificationResult(
            is_creation_knowledge=None,
            label="unsupported_content_mode",
            confidence=None,
            reason="内容模态暂不支持，跳过粗筛",
            prompt_version="content_mode_guard",
            result_payload={"content_mode": content_mode},
            status="skipped",
            error_message="unsupported_content_mode",
        )
    if content_mode == "video_post" and not video_url:
        return ClassificationResult(
            is_creation_knowledge=None,
            label="video_missing",
            confidence=None,
            reason="视频帖缺少可处理的视频地址，跳过粗筛",
            prompt_version="content_mode_guard",
            result_payload={"content_mode": content_mode, "video_url_missing": True},
            status="skipped",
            error_message="video_url_missing",
        )
    if content_mode == "video_post" or (content_mode is None and video_url):
        version = prompt_version("classify_video")
        payload = {
            "platform": platform,
            "title": title,
            "body_text": body_text,
            "video": video_url,
        }
        if trace_writer is None and trace_context is None:
            is_creation, reason, knowledge, points = classify_video(payload, settings)
        else:
            is_creation, reason, knowledge, points = classify_video(
                payload,
                settings,
                trace_writer=trace_writer,
                trace_context=trace_context,
            )
    else:
        version = prompt_version("classify_imgtext")
        payload = {
            "platform": platform,
            "title": title,
            "body_text": body_text,
            "images": image_urls or [],
        }
        if trace_writer is None and trace_context is None:
            is_creation, reason, knowledge, points = classify_imgtext(payload, settings)
        else:
            is_creation, reason, knowledge, points = classify_imgtext(
                payload,
                settings,
                trace_writer=trace_writer,
                trace_context=trace_context,
            )

    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",
    )