"""通过 TikHub 搜索抖音视频，输出 find_agent 统一候选格式。""" from __future__ import annotations import asyncio import json import logging import os import time from typing import Any import httpx from dotenv import load_dotenv from agents.find_agent.support.search_persistence import persist_search_payload from supply_agent.paths import find_project_root logger = logging.getLogger(__name__) DOUYIN_SEARCH_TIKHUB_API = ( "https://api.tikhub.io/api/v1/douyin/search/fetch_video_search_v2" ) DEFAULT_TIMEOUT = 60.0 _MIN_REQUEST_INTERVAL_SECONDS = 1.0 _rate_limit_lock = asyncio.Lock() _last_request_monotonic = 0.0 _env_loaded = False _CONTENT_TYPE_MAP = { "不限": "0", "视频": "1", "图片": "2", "图文": "2", "文章": "3", "0": "0", "1": "1", "2": "2", "3": "3", } _SORT_TYPE_MAP = { "综合排序": "0", "最多点赞": "1", "最新发布": "2", "0": "0", "1": "1", "2": "2", } _PUBLISH_TIME_MAP = { "不限": "0", "一天内": "1", "最近一天": "1", "一周内": "7", "最近一周": "7", "半年内": "180", "最近半年": "180", "0": "0", "1": "1", "7": "7", "180": "180", } _DURATION_MAP = { "不限": "0", "一分钟内": "0-1", "1分钟以内": "0-1", "1-5分钟": "1-5", "五分钟以上": "5-10000", "5分钟以上": "5-10000", "0": "0", "0-1": "0-1", "1-5": "1-5", "5-10000": "5-10000", } def _ensure_env_loaded() -> None: global _env_loaded if _env_loaded: return load_dotenv(find_project_root() / ".env") _env_loaded = True def _safe_int(value: Any, default: int = 0) -> int: if isinstance(value, bool) or value is None: return default try: return int(float(str(value).strip())) except (TypeError, ValueError): return default def _enum_value(value: str, mapping: dict[str, str], field: str) -> str: normalized = str(value).strip() if normalized not in mapping: choices = " / ".join(key for key in mapping if not key.isdigit()) raise ValueError(f"{field} 不支持「{value}」，可选：{choices}") return mapping[normalized] def _get_aweme_info(item: Any) -> dict[str, Any]: if not isinstance(item, dict): return {} data = item.get("data") if not isinstance(data, dict): return {} aweme_info = data.get("aweme_info") return aweme_info if isinstance(aweme_info, dict) else {} def _extract_topics(aweme: dict[str, Any]) -> list[str]: topics: list[str] = [] for item in aweme.get("topic_list") or []: if isinstance(item, str): topics.append(item.strip()) elif isinstance(item, dict): topic = ( item.get("topic_name") or item.get("cha_name") or item.get("hashtag_name") or item.get("name") ) if topic: topics.append(str(topic).strip()) for item in aweme.get("text_extra") or []: if isinstance(item, dict): topic = item.get("hashtag_name") if topic: topics.append(str(topic).strip()) for item in aweme.get("cha_list") or []: if isinstance(item, dict): topic = item.get("cha_name") if topic: topics.append(str(topic).strip()) return list(dict.fromkeys(topic for topic in topics if topic)) def _normalize_aweme(aweme: dict[str, Any]) -> dict[str, Any] | None: aweme_id = str(aweme.get("aweme_id") or "").strip() if not aweme_id: return None author = aweme.get("author") if isinstance(aweme.get("author"), dict) else {} stats = ( aweme.get("statistics") if isinstance(aweme.get("statistics"), dict) else {} ) return { "aweme_id": aweme_id, "desc": str( aweme.get("desc") or aweme.get("item_title") or "无标题" )[:200], "url": f"https://www.douyin.com/video/{aweme_id}", "author": { "nickname": str(author.get("nickname") or "未知作者"), "sec_uid": str(author.get("sec_uid") or ""), }, "statistics": { "digg_count": _safe_int(stats.get("digg_count")), "comment_count": _safe_int(stats.get("comment_count")), "share_count": _safe_int(stats.get("share_count")), "collect_count": _safe_int(stats.get("collect_count")), "play_count": _safe_int(stats.get("play_count")), }, "duration_ms": _safe_int(aweme.get("duration")), "topics": _extract_topics(aweme), } def _summary( keyword: str, results: list[dict[str, Any]], *, filtered_count: int, has_more: bool, next_cursor: int, search_id: str, ) -> str: lines = [ f"TikHub 搜索关键词「{keyword}」", ( f"保留 {len(results)} 条" + (f"，过滤短视频 {filtered_count} 条" if filtered_count else "") + ( f"，还有更多（cursor={next_cursor}, search_id={search_id}）" if has_more else "" ) ), "", ] for index, item in enumerate(results, 1): stats = item["statistics"] lines.extend( [ f"{index}. {item['desc'][:50]}", f" ID: {item['aweme_id']}", f" 链接: {item['url']}", ( f" 作者: {item['author']['nickname']} | " f"sec_uid: {item['author']['sec_uid']}" ), ( f" 数据: 点赞 {stats['digg_count']:,} | " f"评论 {stats['comment_count']:,} | " f"分享 {stats['share_count']:,} | " f"收藏 {stats['collect_count']:,}" ), f" 标签: {'、'.join(item['topics']) or '无'}", "", ] ) return "\n".join(lines).rstrip() def _error_result(error: str) -> str: return json.dumps( {"error": error, "title": "TikHub 抖音搜索失败"}, ensure_ascii=False, ) async def _wait_rate_limit() -> None: global _last_request_monotonic async with _rate_limit_lock: elapsed = time.monotonic() - _last_request_monotonic if elapsed < _MIN_REQUEST_INTERVAL_SECONDS: await asyncio.sleep(_MIN_REQUEST_INTERVAL_SECONDS - elapsed) _last_request_monotonic = time.monotonic() async def _douyin_search_tikhub_raw( keyword: str, content_type: str = "视频", sort_type: str = "综合排序", publish_time: str = "不限", cursor: int = 0, filter_duration: str = "不限", search_id: str = "", backtrace: str = "", min_duration_seconds: int = 0, timeout: float | None = None, ) -> str: """ 使用 TikHub 搜索抖音视频，支持多关键词探索和完整分页状态。 这是 douyin_search 的独立搜索来源。首次搜索 cursor=0、search_id/backtrace 为空； 翻页时必须把上次返回的 next_cursor、search_id、backtrace 原样传回。 Args: keyword: Agent 自主确定的实际搜索词。 content_type: 不限 / 视频 / 图片 / 文章，默认视频；也兼容 TikHub 数字代码。 sort_type: 综合排序 / 最多点赞 / 最新发布；也兼容 0 / 1 / 2。 publish_time: 不限 / 一天内 / 一周内 / 半年内；也兼容 0 / 1 / 7 / 180。 cursor: 首次为 0，翻页使用上次返回的 next_cursor。 filter_duration: 不限 / 一分钟内 / 1-5分钟 / 5分钟以上。 search_id: 翻页状态，必须使用同一搜索返回值。 backtrace: 翻页回溯状态，必须使用同一搜索返回值。 min_duration_seconds: 客户端最短时长过滤，默认 0 表示不过滤。 timeout: 请求超时秒数，默认 60。 Returns: JSON 字符串。search_results 与 douyin_search 格式兼容，并额外包含 duration_ms、topics、收藏数和播放数；分页字段为 has_more、next_cursor、 search_id、backtrace。 """ keyword_text = str(keyword).strip() if not keyword_text: return _error_result("keyword 不能为空") try: content_type_value = _enum_value(content_type, _CONTENT_TYPE_MAP, "content_type") sort_type_value = _enum_value(sort_type, _SORT_TYPE_MAP, "sort_type") publish_time_value = _enum_value( publish_time, _PUBLISH_TIME_MAP, "publish_time" ) duration_value = _enum_value( filter_duration, _DURATION_MAP, "filter_duration" ) except ValueError as exc: return _error_result(str(exc)) _ensure_env_loaded() api_key = os.getenv("TIKHUB_API_KEY", "").strip() if not api_key: return _error_result("未设置环境变量 TIKHUB_API_KEY") start_time = time.time() request_timeout = timeout if timeout is not None else DEFAULT_TIMEOUT payload = { "keyword": keyword_text, "cursor": max(0, int(cursor)), "sort_type": sort_type_value, "publish_time": publish_time_value, "filter_duration": duration_value, "content_type": content_type_value, "search_id": str(search_id or ""), "backtrace": str(backtrace or ""), } try: await _wait_rate_limit() async with httpx.AsyncClient( timeout=request_timeout, trust_env=False, headers={ "Content-Type": "application/json", "Authorization": f"Bearer {api_key}", }, ) as client: response = await client.post(DOUYIN_SEARCH_TIKHUB_API, json=payload) response.raise_for_status() body = response.json() data = body.get("data") if isinstance(body.get("data"), dict) else {} business_data = ( data.get("business_data") if isinstance(data.get("business_data"), list) else [] ) config = ( data.get("business_config") if isinstance(data.get("business_config"), dict) else {} ) next_page = ( config.get("next_page") if isinstance(config.get("next_page"), dict) else {} ) results: list[dict[str, Any]] = [] seen: set[str] = set() filtered_count = 0 minimum_ms = max(0, int(min_duration_seconds)) * 1000 for raw_item in business_data: normalized = _normalize_aweme(_get_aweme_info(raw_item)) if normalized is None or normalized["aweme_id"] in seen: continue if ( minimum_ms and normalized["duration_ms"] and normalized["duration_ms"] < minimum_ms ): filtered_count += 1 continue seen.add(normalized["aweme_id"]) results.append(normalized) has_more = bool(config.get("has_more") in (1, True, "1")) next_cursor = _safe_int(next_page.get("cursor")) next_search_id = str(next_page.get("search_id") or search_id or "") next_backtrace = str( next_page.get("backtrace") or config.get("backtrace") or backtrace or "" ) duration_ms = int((time.time() - start_time) * 1000) result = { "title": f"TikHub 抖音搜索: {keyword_text}", "output": _summary( keyword_text, results, filtered_count=filtered_count, has_more=has_more, next_cursor=next_cursor, search_id=next_search_id, ), "provider": "tikhub", "keyword": keyword_text, "request_params": payload, "results_count": len(results), "filtered_count": filtered_count, "has_more": has_more, "next_cursor": next_cursor, "search_id": next_search_id, "backtrace": next_backtrace, "search_results": results, "duration_ms": duration_ms, } logger.info( "douyin_search_tikhub completed: keyword=%s results=%d has_more=%s duration_ms=%d", keyword_text, len(results), has_more, duration_ms, ) return json.dumps(result, ensure_ascii=False) except httpx.HTTPStatusError as exc: text = exc.response.text[:1000] logger.error( "douyin_search_tikhub HTTP error: keyword=%s status=%d", keyword_text, exc.response.status_code, ) return _error_result(f"HTTP {exc.response.status_code}: {text}") except httpx.TimeoutException: return _error_result(f"请求超时（{request_timeout}秒）") except httpx.RequestError as exc: return _error_result(f"网络错误: {exc}") except Exception as exc: logger.error( "douyin_search_tikhub unexpected error: keyword=%s error=%s", keyword_text, exc, exc_info=True, ) return _error_result(f"未知错误: {exc}") async def douyin_search_tikhub( run_id: str, keyword: str, query_reason: str, source_type: str, source_value: str | None = None, parent_search_id: int | None = None, page_no: int = 1, content_type: str = "视频", sort_type: str = "综合排序", publish_time: str = "不限", cursor: int = 0, filter_duration: str = "不限", search_id: str = "", backtrace: str = "", min_duration_seconds: int = 0, timeout: float | None = None, ) -> str: """ 使用 TikHub 搜索一页抖音视频，返回候选基础信息、数据库 ID 和分页状态。 """ result = await _douyin_search_tikhub_raw( keyword=keyword, content_type=content_type, sort_type=sort_type, publish_time=publish_time, cursor=cursor, filter_duration=filter_duration, search_id=search_id, backtrace=backtrace, min_duration_seconds=min_duration_seconds, timeout=timeout, ) return await asyncio.to_thread( persist_search_payload, result, run_id=run_id, keyword=keyword, query_reason=query_reason, source_type=source_type, source_value=source_value, parent_search_id=parent_search_id, cursor=str(cursor), page_no=page_no, provider="tikhub", content_type=content_type, sort_type=sort_type, publish_time=publish_time, provider_state={"search_id": search_id, "backtrace": backtrace}, )