SupplyAgent/agents/find_agent/tools/douyin_search.py

"""
抖音关键词搜索工具

调用内部爬虫服务进行抖音关键词搜索。
"""
from __future__ import annotations

import asyncio
import json
import logging
import time
from typing import Any, Optional

import httpx

from supply_agent.tools import tool

logger = logging.getLogger(__name__)

_MIN_REQUEST_INTERVAL_SECONDS = 10.1
_rate_limit_lock = asyncio.Lock()
_last_request_monotonic: float = 0.0

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


def _build_search_results(items: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """将 API 原始条目转换为结构化搜索结果。"""
    results = []
    for item in items:
        author = item.get("author", {}) if isinstance(item.get("author"), dict) else {}
        stats = item.get("statistics", {}) if isinstance(item.get("statistics"), dict) else {}
        aweme_id = item.get("aweme_id", "")
        results.append(
            {
                "aweme_id": aweme_id,
                "desc": (item.get("desc") or item.get("item_title") or "无标题")[:100],
                "url": f"https://www.douyin.com/video/{aweme_id}" if aweme_id else "",
                "author": {
                    "nickname": author.get("nickname", "未知作者"),
                    "sec_uid": author.get("sec_uid", ""),
                },
                "statistics": {
                    "digg_count": stats.get("digg_count", 0),
                    "comment_count": stats.get("comment_count", 0),
                    "share_count": stats.get("share_count", 0),
                },
            }
        )
    return results


def _build_output_summary(
    keyword: str,
    items: list[dict[str, Any]],
    has_more: bool,
    cursor_value: str,
) -> str:
    """生成给 LLM 阅读的文本摘要。"""
    lines = [f"搜索关键词「{keyword}」"]
    lines.append(
        f"找到 {len(items)} 条结果"
        + (f"，还有更多（cursor={cursor_value}）" if has_more else "")
    )
    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", {}) if isinstance(item.get("author"), dict) else {}
        author_name = author.get("nickname", "未知作者")
        author_id = author.get("sec_uid", "")

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

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

    return "\n".join(lines)


def _success_result(
    keyword: str,
    data: dict[str, Any],
    items: list[dict[str, Any]],
    has_more: bool,
    cursor_value: str,
    duration_ms: int,
) -> str:
    """构建成功时的 JSON 字符串返回值。"""
    search_results = _build_search_results(items)
    payload = {
        "title": f"抖音搜索: {keyword}",
        "output": _build_output_summary(keyword, items, has_more, cursor_value),
        "keyword": keyword,
        "results_count": len(items),
        "has_more": has_more,
        "next_cursor": cursor_value,
        "search_results": search_results,
        "duration_ms": duration_ms,
    }
    return json.dumps(payload, ensure_ascii=False)


def _error_result(error: str, *, title: str = "抖音搜索失败") -> str:
    """构建失败时的 JSON 字符串返回值。"""
    return json.dumps({"error": error, "title": title}, ensure_ascii=False)


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

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

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

    Returns:
        JSON 字符串，包含 output（文本摘要）和 search_results（结构化列表）。
        search_results 中每项含 aweme_id、desc、author、statistics。
        使用 next_cursor 可获取下一页。
    """
    start_time = time.time()
    request_timeout = timeout if timeout is not None else DEFAULT_TIMEOUT

    try:
        global _last_request_monotonic
        async with _rate_limit_lock:
            now_mono = time.monotonic()
            wait_seconds = _MIN_REQUEST_INTERVAL_SECONDS - (now_mono - _last_request_monotonic)
            if wait_seconds > 0:
                await asyncio.sleep(wait_seconds)
            _last_request_monotonic = time.monotonic()

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

        async with httpx.AsyncClient(timeout=request_timeout) as client:
            response = await client.post(
                DOUYIN_SEARCH_API,
                json=payload,
                headers={"Content-Type": "application/json"},
            )
            response.raise_for_status()
            data = response.json()

        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 = bool(data_block.get("has_more", False))
        cursor_value = str(data_block.get("next_cursor", ""))

        duration_ms = int((time.time() - start_time) * 1000)
        logger.info(
            "douyin_search completed: keyword=%s results=%d has_more=%s duration_ms=%d",
            keyword,
            len(items),
            has_more,
            duration_ms,
        )

        return _success_result(keyword, data, items, has_more, cursor_value, duration_ms)

    except httpx.HTTPStatusError as e:
        logger.error(
            "douyin_search HTTP error: keyword=%s status=%d",
            keyword,
            e.response.status_code,
        )
        return _error_result(f"HTTP {e.response.status_code}: {e.response.text}")
    except httpx.TimeoutException:
        logger.error("douyin_search timeout: keyword=%s timeout=%s", keyword, request_timeout)
        return _error_result(f"请求超时（{request_timeout}秒）")
    except httpx.RequestError as e:
        logger.error("douyin_search network error: keyword=%s error=%s", keyword, e)
        return _error_result(f"网络错误: {e}")
    except Exception as e:
        logger.error("douyin_search unexpected error: keyword=%s error=%s", keyword, e, exc_info=True)
        return _error_result(f"未知错误: {e}")


async def main() -> None:
    result_json = await douyin_search(keyword="养老政策", account_id=DOUYIN_ACCOUNT_ID)
    result = json.loads(result_json)
    if "error" in result:
        print(f"搜索失败: {result['error']}")
    else:
        print(result["output"])
        print(f"\n共 {result['results_count']} 条结果")


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