Agent/examples/content_finder/tools/douyin_search.py

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

调用内部爬虫服务进行抖音关键词搜索。
"""
import asyncio
import json
from datetime import datetime
from typing import Optional, Any, Dict, List

import httpx

from agent.tools import tool, ToolResult


# 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 = "",
    timeout: Optional[float] = None,
) -> ToolResult:
    """
    抖音关键词搜索

    Args:
        keyword: 搜索关键词
        content_type: 内容类型（默认 "视频"）
        sort_type: 排序方式（默认 "综合排序"）
        publish_time: 发布时间范围（默认 "不限"）
        cursor: 分页游标（默认 "0"）
        account_id: 账号ID（可选）
        timeout: 超时时间（秒），默认 60

    Returns:
        ToolResult: 搜索结果 JSON
    """
    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

        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()
        return ToolResult(
            title=f"抖音搜索: {keyword}",
            output=json.dumps(data, ensure_ascii=False, indent=2),
            long_term_memory=f"Searched Douyin for '{keyword}'"
        )
    except httpx.HTTPStatusError as e:
        return ToolResult(
            title="抖音搜索失败",
            output="",
            error=f"HTTP error {e.response.status_code}: {e.response.text}"
        )
    except Exception as e:
        return ToolResult(
            title="抖音搜索失败",
            output="",
            error=str(e)
        )


#
# async def main():
#     result = await douyin_search(
#         keyword="宁艺卓",
#         account_id = "771431186"
#     )
#     print(result.output)
#
# if __name__ == "__main__":
#     asyncio.run(main())