Agent/examples/auto_put_ad_mini/tools/account_material_strategy.py

"""Account-level creative material source strategy.

The default remains historical material recall. Accounts explicitly configured
as AI generated use generated image assets in module B.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass

logger = logging.getLogger(__name__)

MATERIAL_SOURCE_HISTORY = "history"
MATERIAL_SOURCE_AI_GENERATED = "ai_generated"

_AI_SOURCE_VALUES = {
    "ai",
    "ai_generated",
    "ai生成",
    "ai生成素材",
    "生成素材",
    "新生成素材",
}
_HISTORY_SOURCE_VALUES = {
    "",
    "history",
    "历史",
    "历史素材",
    "历史投放素材",
    "历史已投素材",
}
_TRUE_VALUES = {"1", "true", "yes", "y", "是", "允许", "回退"}


@dataclass(frozen=True)
class AccountMaterialStrategy:
    account_id: int
    material_source: str = MATERIAL_SOURCE_HISTORY
    ai_fallback_to_history: bool = False

    @property
    def use_ai_generated(self) -> bool:
        return self.material_source == MATERIAL_SOURCE_AI_GENERATED


def normalize_material_source(raw: object) -> str:
    text = str(raw or "").strip().lower()
    if text in _AI_SOURCE_VALUES:
        return MATERIAL_SOURCE_AI_GENERATED
    if text in _HISTORY_SOURCE_VALUES:
        return MATERIAL_SOURCE_HISTORY
    raise ValueError(f"未知素材来源:{raw}")


def parse_bool_flag(raw: object) -> bool:
    return str(raw or "").strip().lower() in _TRUE_VALUES


def ensure_account_material_strategy_columns() -> None:
    """Add strategy columns when deploying to an older DB schema."""
    from db.connection import get_connection

    conn = get_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT COLUMN_NAME
                FROM information_schema.COLUMNS
                WHERE TABLE_SCHEMA = DATABASE()
                  AND TABLE_NAME = 'ad_creation_account_config'
                  AND COLUMN_NAME IN ('material_source', 'ai_fallback_to_history')
                """
            )
            existing = {row["COLUMN_NAME"] for row in (cur.fetchall() or [])}
            if "material_source" not in existing:
                cur.execute(
                    """
                    ALTER TABLE ad_creation_account_config
                    ADD COLUMN material_source VARCHAR(50) NOT NULL DEFAULT 'history'
                    COMMENT '创意素材来源:history/ai_generated'
                    AFTER feedback_key
                    """
                )
            if "ai_fallback_to_history" not in existing:
                cur.execute(
                    """
                    ALTER TABLE ad_creation_account_config
                    ADD COLUMN ai_fallback_to_history BOOLEAN NOT NULL DEFAULT FALSE
                    COMMENT 'AI生成失败时是否回退历史素材'
                    AFTER material_source
                    """
                )
        conn.commit()
    finally:
        conn.close()


def load_account_material_strategy(account_id: int) -> AccountMaterialStrategy:
    """Load configured material strategy. Missing config defaults to history."""
    from db.connection import get_connection

    try:
        ensure_account_material_strategy_columns()
    except Exception as e:
        logger.warning(
            "[material_strategy] ensure columns failed account=%d, fallback history: %s",
            account_id, e,
        )
        return AccountMaterialStrategy(account_id=account_id)

    conn = get_connection()
    try:
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT material_source, ai_fallback_to_history
                FROM ad_creation_account_config
                WHERE account_id=%s
                LIMIT 1
                """,
                (account_id,),
            )
            row = cur.fetchone() or {}
    finally:
        conn.close()

    try:
        source = normalize_material_source(row.get("material_source"))
    except ValueError:
        logger.warning(
            "[material_strategy] account=%d material_source=%r 非法,按 history 处理",
            account_id, row.get("material_source"),
        )
        source = MATERIAL_SOURCE_HISTORY
    return AccountMaterialStrategy(
        account_id=account_id,
        material_source=source,
        ai_fallback_to_history=bool(row.get("ai_fallback_to_history")),
    )