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

from creation_knowledge.config import load_env_file
from creation_knowledge.jsonio import extract_json_object, to_bool
from creation_knowledge.models import ExtractedContent, Post
from creation_knowledge.prompts import load_prompt

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


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_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 = 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}]
        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 = extract_json_object(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)