Agent/examples/content_finder/tools/aigc_platform_api.py

"""
AIGC接口调用
调用AIGC接口创建爬取计划，绑定生成计划
"""
import json
import logging
import os
from datetime import datetime
from pathlib import Path
from typing import List, Dict, Union, Tuple, Any

import requests

from agent import ToolResult, tool
from db import update_content_plan_ids

logger = logging.getLogger(__name__)

AIGC_BASE_URL = "https://aigc-api.aiddit.com"
CRAWLER_PLAN_CREATE_URL = f"{AIGC_BASE_URL}/aigc/crawler/plan/save"
GET_PRODUCE_PLAN_DETAIL_BY_ID = f"{AIGC_BASE_URL}/aigc/produce/plan/detail"
PRODUCE_PLAN_SAVE = f"{AIGC_BASE_URL}/aigc/produce/plan/save"
DEFAULT_TOKEN = "8bf14f27fc3a486788f3383452422d72"
DEFAULT_TIMEOUT = 60.0


def _load_output_json(trace_id: str, output_dir: str) -> Dict[str, Any]:
    """Load {output_dir}/{trace_id}/output.json."""
    path = Path(output_dir) / trace_id / "output.json"
    if not path.exists():
        raise FileNotFoundError(f"output.json not found: {path}")
    with path.open("r", encoding="utf-8") as f:
        return json.load(f)


def _extract_content_ids(data: Dict[str, Any]) -> List[str]:
    """Extract aweme_id list from output json."""
    contents = data.get("contents") or []
    if not isinstance(contents, list):
        return []
    content_ids: List[str] = []
    for item in contents:
        if not isinstance(item, dict):
            continue
        aweme_id = item.get("aweme_id")
        if aweme_id is None:
            continue
        aweme_id_str = str(aweme_id).strip()
        if aweme_id_str:
            content_ids.append(aweme_id_str)
    return content_ids


def _get_produce_plan_ids_from_env() -> List[str]:
    """Read AIGC_DEMAND_DOUYIN_CONTENT_PRODUCE_PLAN_ID from env."""
    raw = os.getenv("AIGC_DEMAND_DOUYIN_CONTENT_PRODUCE_PLAN_ID", "").strip()
    if not raw:
        return []
    # 接口需要 List[str]，因此把 env 字段（字符串）包装成 list。
    return [raw]


@tool(description="根据抖音账号ID创建爬取计划")
async def create_crawler_plan_by_douyin_account_id(
        account_id: str,
        sort_type: str = "最新",
        produce_plan_ids: List[str] = []
) -> ToolResult:
    """
     根据抖音账号ID创建爬取计划
     Args:
         account_id: 抖音账号ID
         sort_type: 搜索时的视频排序方式（最新/最热），默认最新
         produce_plan_ids: 爬取计划要绑定的生成计划ID，默认为空列表

     Returns:
         ToolResult: 包含以下内容
             - output: 文本格式的爬取计划创建结果摘要
             - metadata.result: 结构化的爬取计划创建结果
                - crawler_info: 爬取计划信息
                    - crawler_plan_id: 创建的爬取计划ID
                    - crawler_plan_name: 创建的爬取计划名称
                    - sort_type: 排序方式
                - produce_plan_infos: 绑定的生成计划信息
                    - produce_plan_id: 生成计划ID
                    - produce_plan_name: 生成计划名称
                    - is_success: 是否成功, true表示绑定成功，false表示绑定失败
                    - msg: 绑定失败时为错误信息，绑定成功则为“成功”
     Note:
         - 建议从 metadata.result 获取结构化数据，而非解析 output 文本
    """

    # 验证 account_id 格式
    if not account_id or not isinstance(account_id, str):
        logger.error(f"create_crawler_plan_by_douyin_account_id invalid account_id: {account_id}")
        return ToolResult(
            title="根据抖音账号ID创建爬取计划失败",
            output="",
            error="account_id 参数无效：必须是非空字符串",
        )

    if not account_id.startswith("MS4wLjABAAAA"):
        logger.error(f"create_crawler_plan_by_douyin_account_id invalid sec_uid format account_id:{account_id}")
        return ToolResult(
            title="根据抖音账号ID创建爬取计划失败",
            output="",
            error=f"account_id 格式错误：必须以 MS4wLjABAAAA 开头，当前值: {account_id[:min(20, len(account_id))]}...",
        )

    if produce_plan_ids is None:
        produce_plan_ids = []

    dt = datetime.now().strftime("%Y%m%d%h%M%s")
    crawler_plan_name = f"【内容寻找Agent自动创建】{dt}_抖音账号ID爬取计划_{account_id[:min(30, len(account_id))]}"
    params = {
        "accountFilters": [],
        "channel": 2,
        "contentFilters": [],
        "contentModal": 4,
        "crawlerComment": 0,
        "crawlerMode": 4,
        "filterAccountMatchMode": 2,
        "filterContentMatchMode": 2,
        "frequencyType": 1,
        "inputModeValues": [
            account_id
        ],
        "modelValueConfig": {
            "sortType": sort_type
        },
        "name": crawler_plan_name,
        "planType": 2,
        "searchModeValues": [],
        "selectModeValues": [],
        "srtExtractFlag": 1,
        "videoKeyFrameType": 1,
        "voiceExtractFlag": 1
    }

    try:

        summary_lines = [f"抖音账号【{account_id}】创建爬取计划"]

        response_json = post(CRAWLER_PLAN_CREATE_URL, params)
        if response_json.get("code") != 0:
            return ToolResult(
                title="根据抖音账号ID创建爬取计划失败",
                output=response_json.get("msg", "接口异常"),
                error=f"create crawler plan interface error",
            )

        crawler_plan_id = response_json.get("data", {}).get("id", "")
        summary_lines.append(f"爬取计划名称: {crawler_plan_name}")
        summary_lines.append(f"    抖音账号ID: {account_id}")
        summary_lines.append(f"    爬取计划ID: {crawler_plan_id}")
        summary_lines.append(f"    爬取计划排序方式: {sort_type}")
        produce_plan_infos: List[Dict[str, str]] = []
        if produce_plan_ids:
            input_source_info = {
                "contentType": 1,
                "inputSourceType": 2,
                "inputSourceValue": crawler_plan_id,
                "inputSourceLabel": f"原始帖子-视频-抖音-内容添加计划-{crawler_plan_name}",
                "inputSourceModal": 4,
                "inputSourceChannel": 2
            }
            produce_plan_infos, msg = crawler_plan_bind_produce_plan(input_source_info, produce_plan_ids)
            if produce_plan_infos:
                for produce_plan_info in produce_plan_infos:
                    summary_lines.append("    绑定的生成计划列表: ")
                    summary_lines.append(f"        生成计划名称: {produce_plan_info.get('produce_plan_name', '')}")
                    summary_lines.append(f"            生成计划ID: {produce_plan_info.get('produce_plan_id', '')}")
                    summary_lines.append(f"            绑定结果: {'绑定成功' if not produce_plan_info.get('msg') else '绑定失败'}")
                    summary_lines.append(f"            信息: {produce_plan_info.get('msg', '成功')}")

        return ToolResult(
            title="根据抖音账号ID创建爬取计划",
            output="\n".join(summary_lines),
            metadata={
                "result": {
                    "crawler_info": {
                        "crawler_plan_id": crawler_plan_id,
                        "crawler_plan_name": crawler_plan_name,
                        "sort_type": sort_type,
                    },
                    "produce_plan_infos": [
                        {
                            "produce_plan_id": produce_plan_info.get("produce_plan_id", ""),
                            "produce_plan_name": produce_plan_info.get("produce_plan_name", ""),
                            "is_success": "绑定成功" if not produce_plan_info.get("msg") else "绑定失败",
                            "msg": produce_plan_info.get("msg", "成功"),
                        }
                        for produce_plan_info in produce_plan_infos
                    ]
                }
            },
            long_term_memory="Create crawler plan by DouYin Account ID",
        )
    except Exception as e:
        logger.error(f"create douyin account crawler plan error: {str(e)}, account_id: {account_id} ")
        return ToolResult(
            title="根据抖音账号ID创建爬取计划失败",
            output="",
            error=f"创建爬取计划错误:{str(e)}",
        )


@tool(description="根据抖音视频ID创建爬取计划")
async def create_crawler_plan_by_douyin_content_id(
        trace_id: str,
) -> ToolResult:
    """
    根据抖音视频ID创建爬取计划
    Args:
        trace_id: 内容寻找任务 trace_id（用于读取 {output_dir}/{trace_id}/output.json）
    Returns:
             Returns:
         ToolResult: 包含以下内容
             - output: 文本格式的爬取计划创建结果摘要
             - metadata.result: 结构化的爬取计划创建结果
                - crawler_info: 爬取计划信息
                    - crawler_plan_id: 创建的爬取计划ID
                    - crawler_plan_name: 创建的爬取计划名称
                    - content_ids: 抖音视频ID列表
                - produce_plan_infos: 绑定的生成计划信息
                    - produce_plan_id: 生成计划ID
                    - produce_plan_name: 生成计划名称
                    - is_success: 是否成功, true表示绑定成功，false表示绑定失败
                    - msg: 绑定失败时为错误信息，绑定成功则为“成功”
    Note:
        - 建议从 metadata.result 获取结构化数据，而非解析 output 文本
    """
    if not trace_id or not isinstance(trace_id, str):
        logger.error(f"create_crawler_plan_by_douyin_content_id invalid trace_id: {trace_id}")
        return ToolResult(
            title="根据抖音内容创建爬取计划失败",
            output="",
            error="trace_id 参数无效: trace_id 必须是非空字符串",
        )

    output_dir = os.getenv("OUTPUT_DIR", ".cache/output")
    try:
        data = _load_output_json(trace_id=trace_id, output_dir=output_dir)
        content_ids = _extract_content_ids(data)
    except Exception as e:
        msg = f"加载/解析 output.json 失败: {e}"
        logger.error(msg, exc_info=True)
        return ToolResult(
            title="根据抖音内容创建爬取计划失败",
            output="",
            error=msg,
        )

    if not content_ids:
        return ToolResult(
            title="根据抖音内容创建爬取计划失败",
            output="",
            error="未在 output.json.contents 中找到有效 aweme_id",
        )
    if len(content_ids) > 100:
        logger.error(
            "create_crawler_plan_by_douyin_content_id invalid content_ids length. "
            f"content_ids.length: {len(content_ids)}"
        )
        return ToolResult(
            title="根据抖音内容创建爬取计划失败",
            output="",
            error=f"content_ids 长度异常: 期望1~100, 实际{len(content_ids)}",
        )

    produce_plan_ids = _get_produce_plan_ids_from_env()
    dt = datetime.now().strftime("%Y%m%d%h%M%s")
    crawler_plan_name = f"【内容寻找Agent自动创建】抖音视频直接抓取-{dt}-抖音"
    params = {
        "channel": 2,
        "contentModal": 4,
        "crawlerComment": 0,
        "crawlerMode": 5,
        "filterAccountMatchMode": 2,
        "filterContentMatchMode": 2,
        "frequencyType": 2,
        "inputModeValues": content_ids,
        "name": crawler_plan_name,
        "planType": 2,
        "searchModeValues": [],
        "srtExtractFlag": 1,
        "videoKeyFrameType": 1,
        "voiceExtractFlag": 1
    }

    try:
        summary_lines = [f"抖音视频爬取计划"]

        response_json = post(CRAWLER_PLAN_CREATE_URL, params)
        if response_json.get("code") != 0:
            return ToolResult(
                title="根据抖音内容ID创建爬取计划失败",
                output=response_json.get("msg", "接口异常"),
                error=f"create crawler plan interface error",
            )

        crawler_plan_id = response_json.get("data", {}).get("id", "")
        summary_lines.append(f"爬取计划名称: {crawler_plan_name}")
        summary_lines.append(f"    抖音视频IDs: {','.join(content_ids)}")
        summary_lines.append(f"    爬取计划ID: {crawler_plan_id}")
        produce_plan_infos: List[Dict[str, str]] = []
        db_updated_rows = 0
        # 环境里的生成计划 ID（字符串）；与是否执行绑定接口无关，用于写库
        env_produce_plan_id = (produce_plan_ids[0] if produce_plan_ids else "").strip()

        if produce_plan_ids:
            input_source_info = {
                "contentType": 1,
                "inputSourceType": 2,
                "inputSourceValue": crawler_plan_id,
                "inputSourceLabel": f"原始帖子-视频-抖音-内容添加计划-{crawler_plan_name}",
                "inputSourceModal": 4,
                "inputSourceChannel": 2
            }
            produce_plan_infos, msg = crawler_plan_bind_produce_plan(input_source_info, produce_plan_ids)
            if produce_plan_infos:
                for produce_plan_info in produce_plan_infos:
                    summary_lines.append("    绑定的生成计划列表: ")
                    summary_lines.append(f"        生成计划名称: {produce_plan_info.get('produce_plan_name', '')}")
                    summary_lines.append(f"            生成计划ID: {produce_plan_info.get('produce_plan_id', '')}")
                    summary_lines.append(f"            绑定结果: {'绑定成功' if not produce_plan_info.get('msg') else '绑定失败'}")
                    summary_lines.append(f"            信息: {produce_plan_info.get('msg', '成功')}")

        # 爬取计划 id 与生成计划 id 任一存在则写库（不依赖是否已配置 produce_plan_ids 去走绑定）
        if (crawler_plan_id or "").strip() or env_produce_plan_id:
            try:
                db_updated_rows = update_content_plan_ids(
                    trace_id=trace_id,
                    aweme_ids=content_ids,
                    crawler_plan_id=crawler_plan_id or "",
                    produce_plan_id=env_produce_plan_id,
                )
            except Exception as e:
                logger.error(f"update content plan ids failed: {e}", exc_info=True)

        return ToolResult(
            title="根据抖音内容ID创建爬取计划",
            output="\n".join(summary_lines),
            metadata={
                "result": {
                    "crawler_info": {
                        "crawler_plan_id": crawler_plan_id,
                        "crawler_plan_name": crawler_plan_name,
                    },
                    "produce_plan_infos": [
                        {
                            "produce_plan_id": produce_plan_info.get("produce_plan_id", ""),
                            "produce_plan_name": produce_plan_info.get("produce_plan_name", ""),
                            "is_success": "绑定成功" if not produce_plan_info.get("msg") else "绑定失败",
                            "msg": produce_plan_info.get("msg", "成功"),
                        }
                        for produce_plan_info in produce_plan_infos
                    ]
                },
                "db": {"updated_rows": db_updated_rows},
            },
            long_term_memory="Create crawler plan by DouYin Content IDs",
        )
    except Exception as e:
        logger.error(f"create douyin content crawler plan error. content_ids: {content_ids}, error: {str(e)}")
        return ToolResult(
            title="根据抖音内容ID创建爬取计划失败",
            output="",
            error=f"创建爬取计划错误:{str(e)}",
        )


def crawler_plan_bind_produce_plan(
        input_source_info: Dict[str, Any],
        produce_plan_ids: List[str],
) -> Tuple[Union[List[Dict[str, str]], None], str]:
    if not input_source_info or not produce_plan_ids:
        return None, f"input_source_info or produce_plan_ids is invalid"
    input_source_check_key = ["inputSourceModal", "inputSourceChannel", "contentType"]
    try:
        if not isinstance(produce_plan_ids, list):
            return None, f"produce_plan_ids is not list"
        result: List[Dict[str, str]] = []
        for produce_plan_id in produce_plan_ids:
            produce_plan_info = {
                "produce_plan_id": produce_plan_id,
            }
            result.append(produce_plan_info)
            # 获取生成计划详情，msg不为空表示获取失败
            produce_plan_detail_info, msg = find_produce_plan_info_by_id(produce_plan_id)
            if msg:
                produce_plan_info["msg"] = msg
                continue

            produce_plan_info["produce_plan_name"] = produce_plan_detail_info.get("name", "")

            input_source_groups = produce_plan_detail_info.get("inputSourceGroups", [])
            if not input_source_groups:
                produce_plan_info["msg"] = "生成计划没有输入源组"
                continue
            # 查询当前爬取计划要添加到的输入源组下标
            input_source_index = 0
            for i in range(len(input_source_groups)):
                input_source_group = input_source_groups[i]
                if not input_source_group.get("inputSources", []):
                    continue
                first_input_source = input_source_group.get("inputSources")[0]
                if all(input_source_info.get(k, 0) == first_input_source.get(k, -1) for k in input_source_check_key):
                    input_source_index = i
                    break

            # 对应的输入源组添加输入源
            input_source_group = input_source_groups[input_source_index]
            input_source_group.get("inputSources", []).append(input_source_info)

            response_json = post(PRODUCE_PLAN_SAVE, produce_plan_detail_info)
            if response_json.get("code") != 0 or not response_json.get("data", {}):
                produce_plan_info["msg"] = response_json.get("msg", "爬取计划绑定生成计划异常")

        return result, ""
    except Exception as e:
        logger.error(f"crawler_plan_bind_produce_plan error. input_source_info: {json.dumps(input_source_info)}, produce_plan_ids: {produce_plan_ids}, error: {str(e)},")
        return None, str(e)


def find_produce_plan_info_by_id(
        produce_plan_id: str,
) -> Tuple[Union[Dict[str, str], None], str]:
    try:
        if not produce_plan_id or not isinstance(produce_plan_id, str):
            return None, f"非法的produce_plan_id: {produce_plan_id}"

        params = {
            "id": produce_plan_id,
        }
        response_json = post(GET_PRODUCE_PLAN_DETAIL_BY_ID, params)

        if response_json.get("code") != 0 or not response_json.get("data", {}):
            return None, response_json.get("msg", "获取生成计划详情异常")

        return response_json.get("data", {}), ""
    except Exception as e:
        logger.error(f"find_produce_plan_info_by_id error. produce_plan_id: {produce_plan_id}, error: {str(e)},")
        return None, str(e)


def post(url: str, params: Any) -> Dict[str, Any]:
    request = {
        "baseInfo": {
            "token": DEFAULT_TOKEN,
        },
        "params": params
    }
    try:
        logger.info(f"invoke aigc platform. url: {url}, request: {json.dumps(request)}")
        response = requests.post(
            url=url,
            json=request,
            headers={"Content-Type": "application/json"},
            timeout=DEFAULT_TIMEOUT
        )
        response.raise_for_status()
        response_json = response.json()

        logger.info(f"invoke aigc platform. url: {url}, request: {json.dumps(request)}, response: {json.dumps(response_json)}")
        return response_json
    except Exception as e:
        logger.error(f"invoke aigc platform error. url: {url}, request: {json.dumps(request)}, error: {str(e)}")
    return {}