Agent/examples/content_finder/tools/douyin_search.py

"""
抖音关键词搜索工具（示例）

调用内部爬虫服务进行抖音关键词搜索。
"""
import asyncio
import logging
import time
from typing import Optional

import requests

from agent.tools import tool, ToolResult

logger = logging.getLogger(__name__)


# API 基础配置
DOUYIN_SEARCH_API = "http://crawapi.piaoquantv.com/crawler/dou_yin/keyword"
DEFAULT_TIMEOUT = 60.0


@tool(description="通过关键词搜索抖音视频内容")
async def douyin_search(
    keyword: str,
    content_type: str = "视频",
    sort_type: str = "综合排序",
    publish_time: str = "不限",
    cursor: str = "0",
    account_id: str = "771431186",
    timeout: Optional[float] = None,
) -> ToolResult:
    """
    抖音关键词搜索

    通过关键词搜索抖音平台的视频内容，支持多种排序和筛选方式。

    Args:
        keyword: 搜索关键词
        content_type: 内容类型（可选：视频/图文, 默认 "视频"）
        sort_type: 排序方式（可选：综合排序/最新发布/最多点赞,默认 "综合排序"）
        publish_time: 发布时间范围（可选：不限/一天内/一周内/半年内,默认 "不限"）
        cursor: 分页游标，用于获取下一页结果，默认 "0"
        account_id: 账号ID（可选）
        timeout: 超时时间（秒），默认 60

    Returns:
        ToolResult: 包含以下内容：
            - output: 文本格式的搜索结果摘要
            - metadata.search_results: 结构化的搜索结果列表
                - aweme_id: 视频ID
                - desc: 视频描述（最多100字符）
                - author: 作者信息
                    - nickname: 作者昵称
                    - sec_uid: 作者ID（完整，约80字符）
                - statistics: 统计数据
                    - digg_count: 点赞数
                    - comment_count: 评论数
                    - share_count: 分享数
            - metadata.raw_data: 原始 API 返回数据

    Note:
        - 使用 cursor 参数可以获取下一页结果
        - 建议从 metadata.search_results 获取结构化数据，而非解析 output 文本
        - author.sec_uid 约 80 字符，使用时不要截断
        - 返回的 cursor 值可用于下一次搜索的 cursor 参数
    """
    start_time = time.time()

    try:
        payload = {
            "keyword": keyword,
            "content_type": content_type,
            "sort_type": sort_type,
            "publish_time": publish_time,
            "cursor": cursor,
            "account_id": account_id
        }

        request_timeout = timeout if timeout is not None else DEFAULT_TIMEOUT

        response = requests.post(
            DOUYIN_SEARCH_API,
            json=payload,
            headers={"Content-Type": "application/json"},
            timeout=request_timeout
        )
        response.raise_for_status()
        data = response.json()

        # 格式化输出摘要
        summary_lines = [f"搜索关键词「{keyword}」"]

        data_block = data.get("data", {}) if isinstance(data.get("data"), dict) else {}
        items = data_block.get("data", []) if isinstance(data_block.get("data"), list) else []
        has_more = data_block.get("has_more", False)
        cursor_value = data_block.get("next_cursor", "")

        summary_lines.append(f"找到 {len(items)} 条结果" + (f"，还有更多（cursor={cursor_value}）" if has_more else ""))
        summary_lines.append("")

        for i, item in enumerate(items, 1):
            aweme_id = item.get("aweme_id", "unknown")
            desc = (item.get("desc") or item.get("item_title") or "无标题")[:50]

            author = item.get("author", {})
            author_name = author.get("nickname", "未知作者")
            author_id = author.get("sec_uid", "")

            stats = item.get("statistics", {})
            digg_count = stats.get("digg_count", 0)
            comment_count = stats.get("comment_count", 0)
            share_count = stats.get("share_count", 0)

            summary_lines.append(f"{i}. {desc}")
            summary_lines.append(f"   ID: {aweme_id}")
            summary_lines.append(f"   链接: https://www.douyin.com/video/{aweme_id}")
            summary_lines.append(f"   作者: {author_name}")
            summary_lines.append(f"   sec_uid: {author_id}")
            summary_lines.append(f"   数据: 点赞 {digg_count:,} | 评论 {comment_count:,} | 分享 {share_count:,}")
            summary_lines.append("")

        duration_ms = int((time.time() - start_time) * 1000)
        logger.info(
            "douyin_search completed",
            extra={
                "keyword": keyword,
                "results_count": len(items),
                "has_more": has_more,
                "cursor": cursor_value,
                "duration_ms": duration_ms
            }
        )

        return ToolResult(
            title=f"抖音搜索: {keyword}",
            output="\n".join(summary_lines),
            long_term_memory=f"Searched Douyin for '{keyword}', found {len(items)} results",
            metadata={
                "raw_data": data,
                "search_results": [  # 结构化搜索结果，供 Agent 直接引用
                    {
                        "aweme_id": item.get("aweme_id"),
                        "desc": (item.get("desc") or item.get("item_title") or "无标题")[:100],
                        "author": {
                            "nickname": item.get("author", {}).get("nickname", "未知作者"),
                            "sec_uid": item.get("author", {}).get("sec_uid", ""),
                        },
                        "statistics": {
                            "digg_count": item.get("statistics", {}).get("digg_count", 0),
                            "comment_count": item.get("statistics", {}).get("comment_count", 0),
                            "share_count": item.get("statistics", {}).get("share_count", 0),
                        }
                    }
                    for item in items
                ]
            }
        )
    except requests.exceptions.HTTPError as e:
        logger.error(
            "douyin_search HTTP error",
            extra={
                "keyword": keyword,
                "status_code": e.response.status_code,
                "error": str(e)
            }
        )
        return ToolResult(
            title="抖音搜索失败",
            output="",
            error=f"HTTP {e.response.status_code}: {e.response.text}"
        )
    except requests.exceptions.Timeout:
        logger.error("douyin_search timeout", extra={"keyword": keyword, "timeout": request_timeout})
        return ToolResult(
            title="抖音搜索失败",
            output="",
            error=f"请求超时（{request_timeout}秒）"
        )
    except requests.exceptions.RequestException as e:
        logger.error("douyin_search network error", extra={"keyword": keyword, "error": str(e)})
        return ToolResult(
            title="抖音搜索失败",
            output="",
            error=f"网络错误: {str(e)}"
        )
    except Exception as e:
        logger.error("douyin_search unexpected error", extra={"keyword": keyword, "error": str(e)}, exc_info=True)
        return ToolResult(
            title="抖音搜索失败",
            output="",
            error=f"未知错误: {str(e)}"
        )


async def main():
    result = await douyin_search(
        keyword="养老政策",
        account_id="771431186"
    )
    print(result.output)

if __name__ == "__main__":
    asyncio.run(main())