xiaoxia-saas/packages/application/viral_video/prompt_loader.py

"""Prompt 模板加载器：从 viral_video_prompt_templates 读模板，30 秒 TTL 热加载。

DB 不可用或没有数据时自动回落到 prompts.DEFAULT_TEMPLATES，保证流程不阻断。
"""

from __future__ import annotations

import threading
import time
from dataclasses import dataclass
from typing import Optional

import sqlalchemy as sa

from packages.adapters.sqlalchemy_impl import session as _session_mod
from packages.application.viral_video.prompts import DEFAULT_TEMPLATES

CACHE_TTL_SECONDS = 30.0

_VALID_TYPES = {"image_analysis", "intent_parsing", "copy_fusion", "storyboard", "review"}


@dataclass
class PromptTemplate:
    name: str
    prompt_type: str
    version: int
    system_prompt: str
    user_prompt_template: str
    example_output: str = ""
    is_active: bool = True


_lock = threading.Lock()
_cache: dict[str, tuple[float, PromptTemplate]] = {}


def _fallback(prompt_type: str) -> Optional[PromptTemplate]:
    for item in DEFAULT_TEMPLATES:
        if item["prompt_type"] == prompt_type:
            return PromptTemplate(
                name=item["name"],
                prompt_type=item["prompt_type"],
                version=item["version"],
                system_prompt=item["system_prompt"],
                user_prompt_template=item["user_prompt_template"],
                example_output=item["example_output"] or "",
                is_active=bool(item["is_active"]),
            )
    return None


_lazy_session = None


def _get_session():
    """优先用全局 SessionLocal（worker）；否则按应用配置懒建同步引擎（api）。"""
    global _lazy_session
    if _session_mod.SessionLocal is not None:
        return _session_mod.SessionLocal()
    if _lazy_session is not None:
        return _lazy_session()
    try:
        from packages.config import get_shared_settings

        url = str(get_shared_settings().database_url)
    except Exception:  # noqa: BLE001
        return None
    if not url:
        return None
    url = url.replace("postgresql+asyncpg://", "postgresql+psycopg://")
    url = url.replace("postgresql://", "postgresql+psycopg://") if url.startswith("postgresql://") else url
    engine = sa.create_engine(url, pool_pre_ping=True, pool_size=2, max_overflow=2)
    from sqlalchemy.orm import sessionmaker

    _lazy_session = sessionmaker(bind=engine)
    return _lazy_session()


def _load_from_db(prompt_type: str) -> Optional[PromptTemplate]:
    session = None
    try:
        session = _get_session()
        if session is None:
            return None
        sql = sa.text("""
            SELECT name, prompt_type, version, system_prompt,
                   user_prompt_template, COALESCE(example_output, '') AS example_output,
                   is_active
              FROM viral_video_prompt_templates
             WHERE prompt_type = :pt AND is_active = TRUE
             ORDER BY version DESC
             LIMIT 1
            """)
        row = session.execute(sql, {"pt": prompt_type}).first()
        if row is None:
            return None
        return PromptTemplate(
            name=row[0],
            prompt_type=row[1],
            version=int(row[2]),
            system_prompt=row[3],
            user_prompt_template=row[4],
            example_output=row[5] or "",
            is_active=bool(row[6]),
        )
    except Exception:  # noqa: BLE001 - 表不存在/DB 不可用时静默回落
        return None
    finally:
        if session is not None:
            try:
                session.close()
            except Exception:  # noqa: BLE001
                pass


def get_template(prompt_type: str, *, force_refresh: bool = False) -> Optional[PromptTemplate]:
    """取某类型当前启用模板，30 秒缓存；DB 无数据则回落到代码默认模板。"""
    if prompt_type not in _VALID_TYPES:
        raise ValueError(f"未知 prompt_type: {prompt_type}")

    now = time.monotonic()
    with _lock:
        cached = _cache.get(prompt_type)
        if not force_refresh and cached and now - cached[0] < CACHE_TTL_SECONDS:
            return cached[1]

    template = _load_from_db(prompt_type) or _fallback(prompt_type)
    if template is not None:
        with _lock:
            _cache[prompt_type] = (now, template)
    return template


def invalidate() -> None:
    """清空缓存（测试用）。"""
    with _lock:
        _cache.clear()


class _SafeDict(dict):
    def __missing__(self, key: str) -> str:
        return "{" + key + "}"


def _safe_format(text: str, kwargs: dict) -> str:
    try:
        return text.format_map(_SafeDict(kwargs))
    except Exception:  # noqa: BLE001
        return text


def render_user_prompt(template: PromptTemplate, **kwargs) -> str:
    """填充 user_prompt_template 占位符，缺键原样保留不报错。"""
    return _safe_format(template.user_prompt_template, kwargs)


def render_system_prompt(template: PromptTemplate, **kwargs) -> str:
    """copy_fusion 等 system_prompt 含运行时变量时填充。"""
    return _safe_format(template.system_prompt, kwargs)