Create-knowledge-find-decode/creation_knowledge/integrations/extractor.py

"""多模态内容理解：把帖子的文本 + 图片(+视频) 交给 Gemini，提取真正的创作知识。

走 OpenRouter /chat/completions，模型 google/gemini-3-flash-preview。
消息格式（system + user.content 为 [text, image_url...] 列表）对齐 ContentFindAgentNew
的 gemini_video.GeminiVideoClient，但提示词换成「提取创作知识」，不做相关性审核。

关键：知识常在图片里，不能只读 body_text —— 见 创作知识-重构设计.md。
"""
from __future__ import annotations

from typing import Any, Callable, Mapping, Optional

import httpx

import logging

from core.config import load_env_file
from core.jsonio import extract_json_object, to_bool
from core.models import Card, CardExtract, ExtractedContent, Post
from core.prompts import load_prompt

logger = logging.getLogger(__name__)

DEFAULT_MODEL = "google/gemini-3-flash-preview"
DEFAULT_BASE_URL = "https://openrouter.ai/api/v1"
DEFAULT_TIMEOUT = 90.0
MAX_CARDS = 12  # 一次提取最多送多少张卡片（图/帧），控成本；超出截断并记日志


def _card_label(card: Card) -> str:
    """给模型看的卡片标签，如 【卡片1】 或 【卡片3 · 00:12】。card.index 为 1-based。"""
    if card.kind == "frame" and card.timestamp is not None:
        ts = int(card.timestamp)
        return f"【卡片{card.index} · {ts // 60:02d}:{ts % 60:02d}】"
    return f"【卡片{card.index}】"

_SYSTEM_PROMPT = (
    "你是创作知识提取助手。从给定的小红书帖子（标题、正文、图片、视频）中，"
    "提取真正能指导『如何创作内容』的知识；知识常在图片/视频里而非正文。"
    "忠实提取、不编造。只输出一个 JSON 对象，不要解释或 markdown。"
)


class ExtractorError(RuntimeError):
    pass


class GeminiExtractor:
    def __init__(
        self,
        *,
        api_key: str,
        model: str = DEFAULT_MODEL,
        base_url: str = DEFAULT_BASE_URL,
        timeout_seconds: float = DEFAULT_TIMEOUT,
        http_post: Callable[..., Any] = httpx.post,
        max_cards: int = MAX_CARDS,
    ) -> None:
        if not api_key:
            raise ExtractorError("missing OPENROUTER_API_KEY")
        self.api_key = api_key
        self.model = model
        self.base_url = base_url.rstrip("/")
        self.timeout_seconds = timeout_seconds
        self.http_post = http_post
        self.max_cards = max_cards

    @classmethod
    def from_env(cls, env: Mapping[str, str] | None = None, env_file: str = ".env") -> "GeminiExtractor":
        source = dict(load_env_file(env_file))
        if env:
            source.update(env)
        api_key = source.get("OPENROUTER_API_KEY") or source.get("OPEN_ROUTER_API_KEY") or ""
        return cls(
            api_key=api_key,
            model=source.get("CONTENT_AGENT_VIDEO_LLM_MODEL") or DEFAULT_MODEL,
            base_url=source.get("OPENROUTER_BASE_URL") or DEFAULT_BASE_URL,
            timeout_seconds=float(source.get("CONTENT_AGENT_VIDEO_LLM_TIMEOUT_SECONDS") or DEFAULT_TIMEOUT),
            max_cards=int(source.get("CK_MAX_CARDS") or MAX_CARDS),
        )

    def _cards(self, post: Post) -> list[Card]:
        """取要送给模型的卡片：优先 post.cards；为空则回退 image_urls。"""
        cards = post.cards or [
            Card(index=i, kind="image", url=u)
            for i, u in enumerate(post.image_urls, start=1)
        ]
        if len(cards) > self.max_cards:
            dropped = [c.index for c in cards[self.max_cards:]]
            logger.warning("post %s 卡片数 %d 超过 MAX_CARDS=%d，截断丢弃卡片 %s",
                           post.id, len(cards), self.max_cards, dropped)
            cards = cards[: self.max_cards]
        return cards

    def build_messages(self, post: Post) -> list[dict]:
        user_text = load_prompt("extract").format(
            title=post.title or "(无)",
            topics="、".join(post.topic_list) or "(无)",
            body=post.body_text or "(空)",
        )
        parts: list[dict] = [{"type": "text", "text": user_text}]
        # 每张卡片前插一个【卡片N】标签，再插图，便于模型按卡片归因
        for card in self._cards(post):
            parts.append({"type": "text", "text": _card_label(card)})
            parts.append({"type": "image_url", "image_url": {"url": card.url}})
        return [
            {"role": "system", "content": _SYSTEM_PROMPT},
            {"role": "user", "content": parts},
        ]

    def extract(self, post: Post) -> ExtractedContent:
        messages = self.build_messages(post)
        last_exc: Optional[Exception] = None
        for attempt in range(2):
            try:
                resp = self.http_post(
                    f"{self.base_url}/chat/completions",
                    headers={
                        "Authorization": f"Bearer {self.api_key}",
                        "Content-Type": "application/json",
                    },
                    json={"model": self.model, "messages": messages},
                    timeout=self.timeout_seconds,
                )
                resp.raise_for_status()
                content = resp.json()["choices"][0]["message"]["content"]
                data = extract_json_object(content)
                cards = []
                for c in data.get("cards") or []:
                    if isinstance(c, dict) and c.get("index") is not None:
                        cards.append(CardExtract(
                            index=int(c["index"]),
                            content=str(c.get("content") or ""),
                        ))
                return ExtractedContent(
                    text=str(data.get("text") or ""),
                    cards=cards,
                    from_image=str(data.get("from_image") or ""),
                    from_video=str(data.get("from_video") or ""),
                    is_empty=to_bool(data.get("is_empty")),
                )
            except httpx.HTTPError as exc:
                last_exc = exc
                if attempt == 0:
                    continue
                raise ExtractorError(f"gemini_http_error: {exc}") from exc
            except (KeyError, IndexError, TypeError, ValueError) as exc:
                last_exc = exc
                if attempt == 0:
                    continue
                raise ExtractorError(f"gemini_response_invalid: {exc}") from exc
        raise ExtractorError(f"gemini_unknown_error: {last_exc}")


def extract_content(
    post: Post,
    *,
    client: Optional[GeminiExtractor] = None,
    env_file: str = ".env",
) -> ExtractedContent:
    client = client or GeminiExtractor.from_env(env_file=env_file)
    return client.extract(post)