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

import json
from typing import Any, Callable, Mapping, Optional

import httpx

from creation_knowledge.config import load_env_file
from creation_knowledge.models import ExtractedContent, Post

DEFAULT_MODEL = "google/gemini-3-flash-preview"
DEFAULT_BASE_URL = "https://openrouter.ai/api/v1"
DEFAULT_TIMEOUT = 90.0
MAX_IMAGES = 6

_SYSTEM_PROMPT = (
    "你是创作知识提取助手。给你一篇小红书帖子（标题、正文、图片），"
    "请提取出其中真正能指导『如何创作内容』的知识。"
    "知识往往在图片里而不在正文里。只输出一个 JSON 对象，不要任何解释或 markdown。"
)

_USER_TMPL = (
    "帖子标题：{title}\n"
    "话题：{topics}\n"
    "正文（可能为空或只是话题串，不要只看它）：\n{body}\n\n"
    "请综合图片内容，输出如下 JSON：\n"
    '{{"text": "把这篇帖子真正讲的创作知识完整、忠实地讲清楚；原文没有的不要编造", '
    '"from_image": "仅来自图片的知识要点；没有则空字符串", '
    '"from_video": "仅来自视频的知识要点；没有则空字符串", '
    '"is_empty": "true 或 false：多模态看完后是否没有任何有效创作知识"}}'
)


class ExtractorError(RuntimeError):
    pass


def _strip_to_json(text: str) -> dict:
    """从模型输出里取出 JSON 对象（容忍 ```json fences 或前后多余文本）。"""
    s = text.strip()
    if s.startswith("```"):
        s = s.split("```", 2)[1] if s.count("```") >= 2 else s.strip("`")
        if s.lstrip().lower().startswith("json"):
            s = s.lstrip()[4:]
    start, end = s.find("{"), s.rfind("}")
    if start == -1 or end == -1 or end < start:
        raise ExtractorError(f"no json object in model output: {text[:120]!r}")
    return json.loads(s[start : end + 1])


def _to_bool(value: Any) -> bool:
    if isinstance(value, bool):
        return value
    return str(value).strip().lower() in ("1", "true", "yes", "是")


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_images: int = MAX_IMAGES,
    ) -> 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_images = max_images

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

    def build_messages(self, post: Post) -> list[dict]:
        user_text = _USER_TMPL.format(
            title=post.title or "(无)",
            topics="、".join(post.topic_list) or "(无)",
            body=post.body_text or "(空)",
        )
        parts: list[dict] = [{"type": "text", "text": user_text}]
        for url in post.image_urls[: self.max_images]:
            parts.append({"type": "image_url", "image_url": {"url": 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 = _strip_to_json(content)
                return ExtractedContent(
                    text=str(data.get("text") or ""),
                    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)