beige/backend/app/workers/pipeline_io.py

"""
app/workers/pipeline_io.py — 生产链 Step5-8
Step5: 文案生成（generate_text_variants）
Step6: 图片生成（generate_storyboard_images，asyncio.gather）
Step7: 图片后处理（image_postprocessor）
Step8: 存 text_candidates / image_candidates → 更新状态 → 推 task_done
"""

import asyncio
import json
import logging
import os

logger = logging.getLogger(__name__)


def _resolve_image_path(img_path: str) -> str:
    """
    解析产品参考图路径，兼容绝对路径(新)与历史相对路径(旧)。
    新数据存绝对路径(/app/uploads/...)直接返回；
    旧数据存相对路径(uploads/packages/...)锚定到 UPLOAD_ABS_ROOT 的父级，
    避免 worker(cwd=/) 解析失败。
    """
    if not img_path:
        return ""
    if os.path.isabs(img_path):
        return img_path
    from app.core.config import get_settings
    # UPLOAD_ABS_ROOT=/app/uploads，其父级 /app 是相对路径(uploads/...)的锚点
    root_parent = os.path.dirname(get_settings().UPLOAD_ABS_ROOT.rstrip("/"))
    return os.path.join(root_parent, img_path)


def run_text_generation(db, clients, task, product_dict: dict, flywheel_fragment: str,
                        push_fn, workspace_id: int, seq_start: int) -> tuple[list, int, bool]:
    """
    Step5: 调 generate_text_variants → 存 TextCandidate → 推 SSE → 写 ai_call_logs。
    S1: 存库前过滤——只存 passed且score>=QUALITY_PASS_SCORE(80,红线)且banned_word_status!='hard_block' 的文案。
    合格数 < task.text_count 时 needs_replenish=True（由主任务发起后台补充子任务）。
    返回 (candidates_raw, next_seq, needs_replenish)。
    """
    import time
    from app.services.ai_engine.text_variants import generate_text_variants
    from app.models.product import BannedWord
    from app.models.task import TextCandidate
    from app.models.flywheel import AiCallLog
    from app.models.workspace import UserApiKey
    from app.constants.enums import CandidateSource, BannedWordStatus
    from app.services.ai_engine.constants import QUALITY_PASS_SCORE

    banned_rows = db.query(BannedWord).filter(
        BannedWord.workspace_id == workspace_id
    ).all()
    banned_dicts = [{"word": b.word, "level": b.level, "replacement": b.replacement}
                    for b in banned_rows]

    # 查 key_id（只取 id，不解密，不违反基石B）
    key_row = db.query(UserApiKey).filter(
        UserApiKey.user_id == task.operator_id,
        UserApiKey.workspace_id == workspace_id,
    ).first()
    key_id = key_row.id if key_row else None

    t0 = time.monotonic()
    llm_success = True
    try:
        candidates_raw = asyncio.run(generate_text_variants(
            llm_client=clients,
            product=product_dict,
            count=task.text_count,
            banned_word_rows=banned_dicts,
            flywheel_context=flywheel_fragment,
        ))
    except Exception as exc:
        llm_success = False
        logger.error("generate_text_variants 失败: %s", exc)
        candidates_raw = []
    latency_ms = int((time.monotonic() - t0) * 1000)

    # 写 ai_call_logs（留痕，不含明文key）
    try:
        log = AiCallLog(
            workspace_id=workspace_id,
            user_id=task.operator_id,
            key_id=key_id,
            task_id=task.id,
            provider="apiports",
            model=clients._model,
            call_type="text",
            success=llm_success,
            latency_ms=latency_ms,
        )
        db.add(log)
        db.flush()
    except Exception as log_exc:
        logger.warning("ai_call_logs 写入失败（非阻断）: %s", log_exc)

    # S1: 存库前过滤——只存合格文案（passed + score>=QUALITY_PASS_SCORE(80) + 非hard_block）
    seq = seq_start
    saved_count = 0
    for i, c in enumerate(candidates_raw):
        score = c.get("score", 0)
        passed = c.get("passed", False)
        bw_status = c.get("banned_word_status", "pass")
        if not (passed and score >= QUALITY_PASS_SCORE and bw_status != "hard_block"):
            logger.info(
                "文案[%d] 过滤丢弃: passed=%s score=%s banned=%s",
                i, passed, score, bw_status,
            )
            continue

        tc = TextCandidate(
            workspace_id=workspace_id,
            task_id=task.id,
            source=CandidateSource.AI,
            angle_label=c.get("angle_label") or c.get("angle", ""),
            content=json.dumps(c, ensure_ascii=False),
            score_json=json.dumps(c.get("score_detail", []), ensure_ascii=False),
            banned_word_status=BannedWordStatus(bw_status),
        )
        db.add(tc)
        db.flush()
        saved_count += 1
        seq += 1
        push_fn(task.id, workspace_id, "text_candidate", {
            "candidate_id": tc.id, "angle_label": tc.angle_label,
            "content": c.get("content", ""), "score": score,
        }, seq)
        seq += 1
        push_fn(task.id, workspace_id, "text_progress", {
            "done": saved_count, "total": task.text_count
        }, seq)

    db.commit()

    # S1: 合格数不足时标记需要后台补充
    needs_replenish = saved_count < task.text_count
    if needs_replenish:
        logger.warning(
            "文案合格数不足: task_id=%s 目标=%s 实得=%s，将后台异步补充",
            task.id, task.text_count, saved_count,
        )
    return candidates_raw, seq, needs_replenish


def run_image_generation(db, clients, task, product_dict: dict,
                         push_fn, workspace_id: int, seq_start: int,
                         first_copy: dict, upload_base_path: str) -> int:
    """
    Step6+7+8(image): 调 generate_storyboard_images → 后处理 → 存 ImageCandidate → 推 SSE。
    返回 next_seq。
    """
    import time
    from app.services.ai_engine.image_gen import generate_storyboard_images
    from app.services.ai_engine.image_postprocessor import process_image
    from app.models.task import ImageCandidate
    from app.models.flywheel import AiCallLog
    from app.models.workspace import UserApiKey
    from app.constants.enums import ImageRole as IR

    # 取 key_id（不解密，不记录明文 key）
    key_row = db.query(UserApiKey).filter(
        UserApiKey.user_id == task.operator_id,
        UserApiKey.workspace_id == workspace_id,
    ).first()
    key_id = key_row.id if key_row else None

    # TODO: 尺寸字段后续加产品级配置（products 表现无 aspect_ratio 字段）
    # 本轮固定 '3:4'=1024×1536，与 gpt-image-2 原生尺寸一致，免后处理二次拉伸
    aspect_ratio = "3:4"

    # image_count=0 直接跳过(纯文案任务/测试)，不空跑生图通道触发无谓失败日志。
    if not task.image_count or task.image_count <= 0:
        logger.info("image_count=0，跳过生图: task_id=%s", task.id)
        return seq_start

    reference_images: list[bytes] = []
    _img_path = _resolve_image_path(product_dict.get("image_path", ""))
    if _img_path and os.path.isfile(_img_path):
        try:
            with open(_img_path, "rb") as _f:
                reference_images = [_f.read()]
            logger.info("产品参考图已加载：%s (%d bytes)", _img_path, len(reference_images[0]))
        except Exception as _e:
            logger.warning("产品参考图读取失败，退化为空列表：%s %s", _img_path, _e)
    else:
        logger.warning(
            "product.image_path 未设置或文件不存在（%r），生图将以无参考图模式运行，"
            "可能导致产品包装跑偏。", _img_path
        )

    # 禁降级兜底：本次产品入镜但无参考图 → 硬失败，绝不降级纯文生图（建任务已拦一道，这是防绕过）
    if getattr(task, "need_product_image", True) and not reference_images:
        raise ValueError(
            "本次产品入镜(need_product_image=True)但未获取到产品参考图，"
            "拒绝降级纯文生图。请确认产品已上传参考图。"
        )

    seq = seq_start
    # 3套正交叙事 A/B/C，每套各 image_count 张独立生图
    for strategy in ("A", "B", "C"):
        t0 = time.monotonic()
        img_success = True
        img_error_code = None
        try:
            image_results = asyncio.run(generate_storyboard_images(
                client=clients,
                note=first_copy,
                product=product_dict,
                image_count=task.image_count,
                reference_images=reference_images or None,
                strategy=strategy,
            ))
        except Exception as exc:
            img_success = False
            img_error_code = type(exc).__name__
            logger.error("generate_storyboard_images 套%s 失败: %s", strategy, exc)
            image_results = []
        latency_ms = int((time.monotonic() - t0) * 1000)

        fail_count = 0
        first_img_error: str | None = None
        for i, img_result in enumerate(image_results):
            if img_result.get("error"):
                fail_count += 1
                if first_img_error is None:
                    first_img_error = str(img_result["error"])[:32]
                seq += 1
                push_fn(task.id, workspace_id, "batch_failed", {
                    "batch": img_result["role"], "reason": img_result["error"],
                    "strategy": strategy, "retryable": True,
                }, seq)
                continue

            raw_bytes = img_result["image_bytes"]
            try:
                processed = process_image(raw_bytes, aspect_ratio=aspect_ratio, resample_strength=1)
            except Exception as e:
                logger.warning("图片后处理失败，使用原图: %s", e)
                processed = raw_bytes

            img_dir = os.path.join(upload_base_path, str(workspace_id), str(task.id))
            os.makedirs(img_dir, exist_ok=True)
            filename = f"{strategy}_{i+1:02d}_{img_result['role']}.jpg"
            img_path = os.path.join(img_dir, filename)
            with open(img_path, "wb") as f:
                f.write(processed)

            img_url = f"/uploads/{workspace_id}/{task.id}/{filename}"

            role_enum = IR.MAIN
            try:
                role_enum = IR(img_result["role"])
            except ValueError:
                pass

            ic = ImageCandidate(
                workspace_id=workspace_id,
                task_id=task.id,
                role=role_enum,
                url=img_url,
                seq=i + 1,
                strategy=strategy,    # 写入 A/B/C（非 hardcode）
            )
            db.add(ic)
            db.flush()
            seq += 1
            push_fn(task.id, workspace_id, "image_candidate", {
                "candidate_id": ic.id, "strategy": strategy,
                "url": img_url, "role": img_result["role"],
            }, seq)
            seq += 1
            push_fn(task.id, workspace_id, "image_progress", {
                "done": i + 1, "total": task.image_count, "strategy": strategy,
            }, seq)

        # 写 ai_call_logs（每套一条，失败不阻断）
        actual_provider = os.environ.get("IMAGE_PROVIDER_PRIMARY", "gpt")
        final_error_code = first_img_error or img_error_code
        try:
            img_log = AiCallLog(
                workspace_id=workspace_id,
                user_id=task.operator_id,
                key_id=key_id,
                task_id=task.id,
                provider=actual_provider,
                call_type="image",
                success=(img_success and fail_count == 0),
                latency_ms=latency_ms,
                error_code=final_error_code,
            )
            db.add(img_log)
            db.flush()
        except Exception as log_exc:
            logger.warning("ai_call_logs(image) 套%s 写入失败（非阻断）: %s", strategy, log_exc)

    db.commit()
    return seq