Create-knowledge-find-decode/creation_knowledge/pipeline.py

"""流水线：把 6 步串起来，每步落库。INGEST_ENABLED=False 时只组装+存，不真实发送。

依赖都可注入（fetch/extract/chat/store），便于离线测试与替换平台。
"""
from __future__ import annotations

import logging
from typing import Callable, Optional

from creation_knowledge import media
from creation_knowledge.config import Settings
from creation_knowledge.integrations.crawler import CrawlerError, fetch_post_detail
from creation_knowledge.integrations.db import CkStore
from creation_knowledge.integrations.extractor import ExtractorError, GeminiExtractor
from creation_knowledge.integrations.llm import ChatFn, default_chat
from creation_knowledge.integrations.video_extract import VideoExtractError, extract_video
from creation_knowledge.ingest import IngestError, ingest as real_ingest
from creation_knowledge.models import ExtractedContent, Post
from creation_knowledge.stages import (
    build_ingest_payload,
    deconstruct_item,
    screen_post,
    split_post,
)

logger = logging.getLogger(__name__)

FetchFn = Callable[[str], Post]
ExtractFn = Callable[[Post], ExtractedContent]


def _skip_reason(post: Post) -> Optional[str]:
    """拿不到可提炼素材的帖子 → 跳过（不报错、不留脏数据）。在提取前判定。

    - video_no_direct_url：声明 video 但无视频直链（小红书视频帖结构性如此，详情只给封面）。
    - empty_media：既无图也无视频（接口降级 / 空帖）。
    """
    if (post.content_type or "").lower() == "video" and not post.video_urls:
        return "video_no_direct_url"
    if not post.image_urls and not post.video_urls:
        return "empty_media"
    return None


def _process_one(
    url: str,
    *,
    settings: Settings,
    store: CkStore,
    fetch_fn: FetchFn,
    extract_fn: ExtractFn,
    chat: ChatFn,
    ingest_enabled: bool,
) -> dict:
    # 1) 拉取
    try:
        post = fetch_fn(url)
    except CrawlerError as exc:
        return {"url": url, "status": "fetch_failed", "error": str(exc)}
    store.upsert_post(post)  # 图文卡片在此落；视频段卡在提取后补落

    # 1.5) 跳过判定：素材够不着的帖子直接 skip（如小红书视频帖）
    skip = _skip_reason(post)
    if skip:
        store.clear_items(post.id)  # 幂等：本帖若是重跑，清掉上轮派生的片段
        store.update_stage(post.id, "skipped")
        return {"url": url, "post_id": post.id, "status": "skipped", "reason": skip}

    # 2) 多模态提取（图文=逐图；视频=原生整段视频→段卡，extract 内写 post.cards）
    try:
        content = extract_fn(post)
    except (ExtractorError, VideoExtractError) as exc:
        store.update_stage(post.id, "failed")
        return {"url": url, "post_id": post.id, "status": "extract_failed", "error": str(exc)}
    if post.cards:
        store.upsert_post(post)  # 视频段卡落库（图文重复 upsert 无害）
    store.set_extracted(post.id, content.model_dump())  # stage=extracted
    store.clear_items(post.id)  # 幂等：重跑先清旧片段，再写本轮（rejected 路径也保持干净）

    # 3) 筛选
    screening = screen_post(post, content, chat=chat)
    store.set_screening(post.id, screening.model_dump())  # stage=screened
    if not screening.passed:
        store.update_stage(post.id, "rejected")
        return {"url": url, "post_id": post.id, "status": "rejected",
                "score": screening.score, "reason": screening.reason}

    # 4) 拆分 -> 5) 解构 -> 6) 组装（+可选入库）
    items = split_post(post, content, chat=chat)
    item_ids = []
    for item in items:
        deco = deconstruct_item(item, chat=chat)
        payload = build_ingest_payload(post, item, deco)
        item_id = store.save_item(
            post.id, item.model_dump(), deco.model_dump(), payload.model_dump()
        )
        item_ids.append(item_id)
        if ingest_enabled:
            try:
                res = real_ingest(payload, settings=settings)
                store.update_item_ingest(item_id, "ingested", res.get("knowledge_id"))
            except IngestError:
                store.update_item_ingest(item_id, "failed", None)
    store.update_stage(post.id, "done")
    return {"url": url, "post_id": post.id, "status": "done", "items": len(item_ids)}


def run_pipeline(
    urls: list[str],
    *,
    settings: Optional[Settings] = None,
    env_file: str = ".env",
    ingest_enabled: Optional[bool] = None,
    run_id: Optional[str] = None,
    store: Optional[CkStore] = None,
    fetch_fn: Optional[FetchFn] = None,
    extract_fn: Optional[ExtractFn] = None,
    chat: Optional[ChatFn] = None,
) -> list[dict]:
    settings = settings or Settings.from_env(env_file)
    ingest_enabled = settings.ingest_enabled if ingest_enabled is None else ingest_enabled
    batch = run_id or media.run_id()
    store = store or CkStore(settings.pg)
    fetch_fn = fetch_fn or (lambda url: fetch_post_detail(url, settings=settings))
    if extract_fn is None:
        _image_client = GeminiExtractor.from_env(env_file=env_file)

        def _dispatch_extract(post: Post) -> ExtractedContent:
            # settings.data_dir 非空 → 媒体落盘到 data/<batch>/<platform>/<post_id>/（解耦开关）
            tgt = (media.media_target(settings.data_dir, batch, post.platform, post.id)
                   if settings.data_dir else None)
            # 视频帖 → 原生整段视频（OpenRouter base64）；图文帖 → 逐图提取
            if (post.content_type or "").lower() == "video" or post.video_urls:
                return extract_video(
                    post, settings=settings,
                    save_path=(tgt.dir / "video.mp4") if tgt else None,
                    public_url=(f"{tgt.public_base}/video.mp4") if tgt else None,
                )
            content = _image_client.extract(post)
            if tgt:  # 图文帖：下载图片到本地并把 card.url 改写为 /data 路径
                media.save_images(post, tgt.dir, tgt.public_base)
            return content

        extract_fn = _dispatch_extract
    chat = chat or default_chat(env_file)

    results = []
    for url in urls:
        results.append(_process_one(
            url, settings=settings, store=store, fetch_fn=fetch_fn,
            extract_fn=extract_fn, chat=chat, ingest_enabled=ingest_enabled,
        ))
    return results