diary-news/backend/app/services/llm/enrichment.py

"""LLM 智能增强服务(翻译后调)。

4 个独立任务:
  1. format   — 排版译文(写入 body_zh_formatted)
  2. classify — 分类(写入 category,多标签)
  3. image    — 生成插图(写入 image_ai_url,prompt 用正文第一段)
  4. commentary — 写点评(写入 commentary)

排版容器 CSS(固定,不再让用户改):
- 字体: system-ui 字体栈
- 字号: 17px
- 行高: 1.7
- 颜色: #3e3e3e
- 段落: margin-bottom 1.5em(自动空一行)

设计:
- 任务入口: enrich_article(article_id, settings_row)
- 任务间互不影响:每个任务独立 try/except + 写 status
- 全部任务共走 LlmClient 的全局限速
- 若设置 enabled=False,只跳过(不调 LLM)
- 用户提示词模板可能不包含全部占位符,用 _safe_format 容错
"""
from __future__ import annotations

import asyncio
import logging
from typing import Any, Mapping

from sqlalchemy import select

from app.database import AsyncSessionLocal
from app.models.article import Article
from app.models.llm_setting import LlmSetting
from app.schemas.llm import get_default_prompts
from app.services.llm.client import LlmClient

logger = logging.getLogger("news.llm.enrichment")

# === 排版容器固定 CSS(项目级固定,不再让用户改)===
# 同时内联到 body_zh_formatted 的容器 div 的 style 属性上,
# 保证分享/邮件/导出场景下样式不丢;前端全局 .article-body 类做兜底。
ARTICLE_BODY_FONT_FAMILY = (
    "system-ui, -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, "
    "'Helvetica Neue', sans-serif"
)
ARTICLE_BODY_FONT_SIZE = "17px"
ARTICLE_BODY_LINE_HEIGHT = "1.7"
ARTICLE_BODY_COLOR = "#3e3e3e"
ARTICLE_BODY_P_MARGIN_BOTTOM = "1.5em"

# === 插图默认尺寸(适中,不再用 1024x768)===
# 写死到 enrichment 里,行为稳定;setting.image_size 仍可由用户在 UI 改,
# 但默认行为不依赖它,避免意外被改成很大。
DEFAULT_IMAGE_SIZE = "768x512"
DEFAULT_IMAGE_FIRST_PARA_CHARS = 400  # 提取第一段最多用这么多字
DEFAULT_IMAGE_MAX_TAGS = 5  # 分类标签上限(多标签)


class _SafeDict(dict):
    """missing 返回 {key} 本身(占位符原样保留),不抛 KeyError。"""

    def __missing__(self, key: str) -> str:  # type: ignore[override]
        return "{" + key + "}"


def _safe_format(template: str, vars_: Mapping[str, Any]) -> str:
    """用 _SafeDict 跑 str.format,缺失的占位符保留原样而不是 KeyError。

    用途:数据库里用户已存的 prompt 模板可能是旧版的(只支持部分占位符),
    新代码传了更多变量也不应崩。
    """
    try:
        return template.format_map(_SafeDict(vars_))
    except (KeyError, IndexError) as e:
        # 极端情况(比如 {} 这种非法占位符)兜底
        logger.warning("_safe_format 解析失败,按原文返回: %s", e)
        return template


# === 获取当前设置(行锁 + 缓存刷新)===
async def get_setting() -> LlmSetting:
    """读 llm_settings 单行;不存在则用默认值插入。"""
    async with AsyncSessionLocal() as session:
        row = (await session.execute(select(LlmSetting).where(LlmSetting.id == 1))).scalar_one_or_none()
        if row is None:
            defaults = get_default_prompts()
            row = LlmSetting(
                id=1,
                format_prompt=defaults["format_prompt"],
                classify_prompt=defaults["classify_prompt"],
                commentary_prompt=defaults["commentary_prompt"],
                image_prompt_template=defaults["image_prompt_template"],
            )
            session.add(row)
            await session.commit()
            await session.refresh(row)
        return row


# === 单任务:format ===
async def _enrich_format(article: Article, setting: LlmSetting, client: LlmClient) -> None:
    template = setting.format_prompt or get_default_prompts()["format_prompt"]
    prompt = _safe_format(template, {"body": (article.body_zh_text or "")[:6000]})
    text = await client.chat(
        system="你是中文新闻排版助手,只输出排版后的纯文本。",
        user=prompt,
        temperature=0.3,
        max_tokens=2000,
    )
    # 极简 HTML 包裹:按段切 + <p>,整体包到带固定 CSS 的 <div> 里
    parts = [f"<p>{p.strip()}</p>" for p in text.split("\n\n") if p.strip()]
    if not parts:
        article.body_zh_formatted = None
    else:
        article.body_zh_formatted = _wrap_article_body("\n".join(parts))
    article.format_status = "ok"


def _wrap_article_body(inner_html: str) -> str:
    """把排版好的段落包到带固定 CSS 的 <div class='article-body'> 里。

    CSS 同时内联到 style 属性(分享/导出样式不丢)+ class 名(前端全局类可覆盖)。
    """
    inline_style = (
        f"font-family:{ARTICLE_BODY_FONT_FAMILY};"
        f"font-size:{ARTICLE_BODY_FONT_SIZE};"
        f"line-height:{ARTICLE_BODY_LINE_HEIGHT};"
        f"color:{ARTICLE_BODY_COLOR};"
    )
    # 段落样式也内联,保证 v-html 渲染时一定生效
    p_style = f"margin:0 0 {ARTICLE_BODY_P_MARGIN_BOTTOM} 0;"
    inner_with_p_style = inner_html.replace("<p>", f'<p style="{p_style}">')
    return f'<div class="article-body" style="{inline_style}">{inner_with_p_style}</div>'


# === 单任务:classify ===
async def _enrich_classify(article: Article, setting: LlmSetting, client: LlmClient) -> None:
    template = setting.classify_prompt or get_default_prompts()["classify_prompt"]
    # 老 prompt 可能只支持 {title}/{summary},不支持 {body} —— _safe_format 兜底
    vars_ = {
        "title": (article.title_zh or article.title)[:200],
        "summary": (article.summary_zh or "")[:400],
        "body": (article.body_zh_text or "")[:1500],
    }
    prompt = _safe_format(template, vars_)
    result = await client.classify_json(
        system="你是新闻分类助手,只返回 JSON。",
        user=prompt,
    )
    cats = result.get("categories") or result.get("tags") or []
    if isinstance(cats, list) and cats:
        # 多标签(2-5 个),逗号分隔存到 category 字段(已有索引)
        joined = ",".join(str(c).strip() for c in cats[:DEFAULT_IMAGE_MAX_TAGS] if str(c).strip())
        article.category = joined[:64] or None
    article.classify_status = "ok"


# === 单任务:image ===
async def _enrich_image(article: Article, setting: LlmSetting, client: LlmClient) -> None:
    template = setting.image_prompt_template or get_default_prompts()["image_prompt_template"]
    # 用正文第一段作为 prompt(英文 prompt 走 title 仍可工作,所以 title 也带上作 fallback)
    first_para = _first_paragraph(article.body_zh_text or "", max_chars=DEFAULT_IMAGE_FIRST_PARA_CHARS)
    if not first_para:
        first_para = (article.title_zh or article.title or "")[:200]
    title_for_prompt = (article.title_zh or article.title or "")[:200]
    # template 同时支持 {body} 和 {title} 两种占位符;老的只支持 {title} 也能跑
    prompt = _safe_format(template, {"body": first_para, "title": title_for_prompt})
    url = await client.generate_image(prompt, size=DEFAULT_IMAGE_SIZE)
    article.image_ai_url = url
    article.image_ai_status = "ok"


def _first_paragraph(text: str, max_chars: int) -> str:
    """取正文第一段(按 \\n\\n 切)。如果首段超长就截断。"""
    if not text:
        return ""
    for p in text.split("\n\n"):
        p = p.strip()
        if p:
            return p[:max_chars]
    return ""


# === 单任务:commentary ===
async def _enrich_commentary(article: Article, setting: LlmSetting, client: LlmClient) -> None:
    template = setting.commentary_prompt or get_default_prompts()["commentary_prompt"]
    prompt = _safe_format(
        template,
        {
            "title": (article.title_zh or article.title)[:200],
            "body": (article.body_zh_text or "")[:3000],
        },
    )
    text = await client.chat(
        system="你是资深新闻评论员。",
        user=prompt,
        temperature=0.6,
        max_tokens=600,
    )
    article.commentary = text or None
    article.commentary_status = "ok"


# === 总编排:enrich_article ===
async def enrich_article(article_id: int) -> dict[str, str]:
    """对单篇文章做 4 项 LLM 增强。

    返回 {task: status} 字典(用于日志)。
    """
    async with AsyncSessionLocal() as session:
        art = (
            await session.execute(select(Article).where(Article.id == article_id))
        ).scalar_one_or_none()
        if not art:
            logger.warning("enrich_article: id=%s not found", article_id)
            return {}
        if not (art.title_zh or art.body_zh_text):
            logger.info("enrich_article: id=%s no translation yet, skip", article_id)
            return {}

    # 拉取设置
    setting = await get_setting()
    if not setting.enabled:
        logger.info("enrich_article: llm disabled, skip id=%s", article_id)
        return {"format": "skipped", "classify": "skipped", "image": "skipped", "commentary": "skipped"}

    # 用配置生成 client(允许热改设置)
    client = LlmClient(
        chat_model=setting.chat_model,
        image_model=setting.image_model,
        interval_sec=setting.interval_sec,
    )

    results: dict[str, str] = {}

    async with AsyncSessionLocal() as session:
        art = (
            await session.execute(select(Article).where(Article.id == article_id))
        ).scalar_one_or_none()
        if not art:
            return {}

        # 4 个任务(互不影响);format / classify / commentary 是 chat,image 是 image
        # 串行执行(已经过 client 内部 Semaphore),但每个 try/except 独立
        tasks: list[tuple[str, Any]] = [
            ("format", _enrich_format(art, setting, client)),
            ("classify", _enrich_classify(art, setting, client)),
            ("image", _enrich_image(art, setting, client)),
            ("commentary", _enrich_commentary(art, setting, client)),
        ]
        for name, coro in tasks:
            try:
                await coro
                results[name] = "ok"
            except Exception as e:
                logger.exception("enrich %s failed for article %s: %s", name, article_id, e)
                results[name] = f"failed:{type(e).__name__}"
                # 标 status
                if name == "format":
                    art.format_status = "failed"
                elif name == "classify":
                    art.classify_status = "failed"
                elif name == "image":
                    art.image_ai_status = "failed"
                elif name == "commentary":
                    art.commentary_status = "failed"

        await session.commit()
    logger.info("enrich_article id=%s: %s", article_id, results)
    return results


# === 后台循环 ===
# 与 translation_loop 一样,常驻从队列里取文章
ENRICHMENT_INTERVAL_SEC = 5.0  # 没活时等待
ENRICHMENT_BATCH_SIZE = 1


async def enrichment_loop() -> None:
    """扫描已翻译但未 enrich 的文章(任一 *_status 为 pending/n/a 且 translation_status=ok)。

    跟 translation_loop 一样常驻。
    """
    logger.info("enrichment_loop started")
    # 等一下让翻译先跑
    await asyncio.sleep(10)
    while True:
        try:
            async with AsyncSessionLocal() as session:
                # 已翻译完成 + 4 个状态中至少有一个是 pending
                rows = (
                    await session.execute(
                        select(Article)
                        .where(
                            Article.translation_status == "ok",
                            Article.title_zh.is_not(None),
                        )
                        .order_by(Article.translated_at.asc().nullslast(), Article.id.asc())
                        .limit(ENRICHMENT_BATCH_SIZE * 5)  # 多取几个找需要 enrich 的
                    )
                ).scalars()
                candidates = list(rows)

            # 过滤:任一 *_status 是 pending
            todo_ids: list[int] = []
            for a in candidates:
                statuses = [
                    a.format_status or "pending",
                    a.classify_status or "pending",
                    a.image_ai_status or "pending",
                    a.commentary_status or "pending",
                ]
                if any(s in ("pending", "failed", "n/a") for s in statuses):
                    todo_ids.append(a.id)
                if len(todo_ids) >= ENRICHMENT_BATCH_SIZE:
                    break

            if not todo_ids:
                await asyncio.sleep(ENRICHMENT_INTERVAL_SEC)
                continue

            for aid in todo_ids:
                try:
                    await enrich_article(aid)
                except Exception as e:
                    logger.exception("enrich_article %s in loop failed: %s", aid, e)
                await asyncio.sleep(0.5)  # 文章间轻节流
        except Exception as e:
            logger.exception("enrichment_loop error: %s", e)
            await asyncio.sleep(ENRICHMENT_INTERVAL_SEC)