beige/backend/app/workers/pipeline_steps.py

"""
app/workers/pipeline_steps.py — 生产链 Step1-4
Step1: 查 DB（task/product）
Step2: 查 key → Fernet 解密（局部变量，不传出，基石B）
Step3: 构建 AI clients
Step4: 推 task_started SSE + 飞轮上下文
"""

import json
import logging

logger = logging.getLogger(__name__)


def load_task_and_product(db, task_id: int):
    """Step1: 查任务 + 产品，失败返 None 或抛异常。"""
    from app.models.task import GenerationTask
    from app.models.product import Product

    task = db.query(GenerationTask).filter(GenerationTask.id == task_id).first()
    if not task:
        logger.error("task_id=%s not found", task_id)
        return None, None

    product = db.query(Product).filter(Product.id == task.product_id).first()
    if not product:
        raise ValueError(f"product_id={task.product_id} not found")
    return task, product


def decrypt_user_key(db, operator_id: int, workspace_id: int) -> str:
    """
    Step2: 查 key → Fernet 解密，返回 plain_key（只活在调用方局部变量）。
    绝不打印、不持久化 plain_key（基石B）。
    """
    from app.models.workspace import UserApiKey
    from app.utils.fernet_utils import decrypt_key

    api_key_row = db.query(UserApiKey).filter(
        UserApiKey.user_id == operator_id,
        UserApiKey.workspace_id == workspace_id,
        UserApiKey.provider.in_(["openai", "apiports"]),  # G6坑修复：接受主备通道名
    ).first()
    if not api_key_row:
        raise ValueError("用户未配置 API Key，请先录入")
    return decrypt_key(api_key_row.encrypted_key)


def decrypt_codeproxy_key(db, operator_id: int, workspace_id: int) -> str | None:
    """
    查用户录入的 codeproxy 备用站 key → Fernet 解密，返回 plain_key 或 None。
    备用通道：用户没录则返回 None（不抛错），build_ai_clients 会回落 env，
    主生图流程绝不因没录备用 key 而中断（apiports 主通道才是必需）。
    """
    from app.models.workspace import UserApiKey
    from app.utils.fernet_utils import decrypt_key

    row = db.query(UserApiKey).filter(
        UserApiKey.user_id == operator_id,
        UserApiKey.workspace_id == workspace_id,
        UserApiKey.provider == "codeproxy",
    ).first()
    return decrypt_key(row.encrypted_key) if row else None


def build_clients_and_clear_key(plain_key: str, alt_key: str | None = None):
    """
    Step3: 构建 AIClients，plain_key 传入后立即由调用方置 None。
    alt_key：用户录入的 codeproxy 备用 key（可选），同样由调用方查库解密后传入。
    返回 clients 对象。
    """
    from app.services.ai_engine.gemini_factory import build_ai_clients
    return build_ai_clients(plain_key, alt_key=alt_key)


def build_product_dict(product) -> dict:
    """把 ORM product 转成 AI 引擎所需的 dict（不含任何 key）。"""
    return {
        "id": getattr(product, "id", None),
        "name": product.name,
        "category": product.category or "通用好物",
        "selling_points": json.loads(product.selling_points or "[]"),
        "style_tone": product.style_tone or "素人分享风",
        "text_angles": json.loads(product.text_angles or "[]"),
        "custom_prompt": product.custom_prompt or "",
        "brand_keyword": product.brand_keyword or "",       # S3: 品牌词透传进生成prompt(每条植入)
        "target_audience": product.target_audience or "",   # 012: 人群透传进storyboard/文案prompt
        "image_path": product.image_path or "",             # 产品参考图路径（主图，向后兼容）
        # R5多图：每张产品图 {path, scene}，生图按分镜role选对应场景图
        "images": [
            {"path": im.path, "scene": im.scene}
            for im in (getattr(product, "images", None) or [])
        ],
        # 第2环标杆配方，默认空；走 AI 主链时由 load_benchmark_features 覆盖填充
        "benchmark_refs": [],
    }


def load_benchmark_features(db, task, workspace_id: int) -> list[dict]:
    """
    第2环→第5环接线：读 task.benchmark_ids → 查 analyze_status=done 的标杆 features_json。
    返回 8维配方 dict 列表（供 build_prompt 借方法层结构，禁抄竞品品牌/功效原话）。
    未选/未分析完/解析失败都安全返空，绝不阻断生成。
    """
    from app.models.product import BenchmarkNote

    raw_ids = getattr(task, "benchmark_ids", None)
    if not raw_ids:
        return []
    try:
        ids = [int(i) for i in json.loads(raw_ids)]
    except Exception:
        return []
    if not ids:
        return []

    rows = db.query(BenchmarkNote).filter(
        BenchmarkNote.id.in_(ids),
        BenchmarkNote.workspace_id == workspace_id,
        BenchmarkNote.analyze_status == "done",
    ).all()

    feats: list[dict] = []
    for b in rows:
        if not b.features_json:
            continue
        try:
            feats.append(json.loads(b.features_json))
        except Exception:
            logger.warning("标杆 features_json 解析失败 id=%s", b.id)
    return feats


def load_flywheel_context(db, workspace_id: int, product_id: int, product_dict: dict) -> tuple[str, dict]:
    """
    查最近50条飞轮事件，聚合偏好上下文。
    返回 (prompt_fragment, full_ctx)。
    """
    from app.models.flywheel import PreferenceEvent
    from app.services.ai_engine.preference_aggregator import aggregate_preference_context

    recent = db.query(PreferenceEvent).filter(
        PreferenceEvent.workspace_id == workspace_id,
        PreferenceEvent.product_id == product_id,
    ).order_by(PreferenceEvent.id.desc()).limit(50).all()

    events_dicts = [
        {"signal_type": e.signal_type, "workspace_id": e.workspace_id,
         "product_id": e.product_id, "angle_label": e.angle_label or "",
         "signal_weight": e.signal_weight, "reason": e.reason or ""}
        for e in recent
    ]
    ctx = aggregate_preference_context(events_dicts, product_dict, workspace_id, product_id)
    # 累积感知：补该产品累计信号总数（前端"飞轮已积累N条信号"），与展示端同口径
    from app.services.flywheel_service import count_signals
    ctx["signal_count"] = count_signals(db, workspace_id, product_id)
    return ctx.get("prompt_fragment", ""), ctx