kg_agent/script/search_recommendations/xiaohongshu_search_recommendations.py

#!/usr/bin/env python3
"""
小红书搜索推荐词接口
获取小红书平台搜索框中的推荐词
"""

import requests
import json
import os
import argparse
import time
import ast
from datetime import datetime
from typing import Dict, Any


class XiaohongshuSearchRecommendations:
    """小红书搜索推荐词API封装类"""

    BASE_URL = "http://47.84.182.56:8001"
    TOOL_NAME = "Xiaohongshu_Search_Recommendations"
    PLATFORM = "xiaohongshu"  # 平台名称

    def __init__(self, results_dir: str = None):
        """
        初始化API客户端

        Args:
            results_dir: 结果输出目录，默认为脚本所在目录下的 results 文件夹
        """
        self.api_url = f"{self.BASE_URL}/tools/call/{self.TOOL_NAME}"

        # 设置结果输出目录
        if results_dir:
            self.results_base_dir = results_dir
        else:
            # 默认使用项目根目录的 data/search_recommendations 文件夹
            script_dir = os.path.dirname(os.path.abspath(__file__))
            project_root = os.path.dirname(os.path.dirname(script_dir))
            self.results_base_dir = os.path.join(project_root, "data", "search_recommendations")

    def get_recommendations(self, keyword: str, timeout: int = 300, max_retries: int = 10, retry_delay: int = 2) -> Dict[str, Any]:
        """
        获取小红书搜索推荐词

        Args:
            keyword: 搜索关键词，例如：'长沙'、'美妆'等
            timeout: 请求超时时间（秒），默认300秒
            max_retries: 最大重试次数，默认10次
            retry_delay: 重试间隔时间（秒），默认2秒

        Returns:
            API响应的JSON数据

        Raises:
            requests.exceptions.RequestException: 请求失败时抛出异常
        """
        payload = {"keyword": keyword}
        last_error = None

        for attempt in range(max_retries + 1):
            try:
                response = requests.post(
                    self.api_url,
                    json=payload,
                    timeout=timeout,
                    headers={"Content-Type": "application/json"}
                )
                response.raise_for_status()
                res = response.json()
                # 使用 ast.literal_eval 解析 Python 字典字符串（不是标准 JSON）
                result = ast.literal_eval(res['result'])

                # 成功：code == 0
                if result.get('code') == 0:
                    return result['data']['data']

                # 失败：code != 0
                last_error = f"code={result.get('code')}"

            except Exception as e:
                from traceback import print_exc
                print(f"发生异常: {e}")
                print_exc()
                last_error = str(e)

            # 统一处理重试逻辑
            if attempt < max_retries:
                print(f"请求失败 ({last_error}), 第{attempt + 1}次重试，等待{retry_delay}秒...")
                time.sleep(retry_delay)
            else:
                print(f"达到最大重试次数({max_retries})，最后错误: {last_error}")

        return []

    def save_result(self, keyword: str, result: Dict[str, Any]) -> str:
        """
        保存结果到文件
        目录结构: results/平台/关键词/时间戳.json

        Args:
            keyword: 搜索关键词
            result: API返回的结果

        Returns:
            保存的文件路径
        """
        # 创建目录结构: results/平台/关键词/
        result_dir = os.path.join(self.results_base_dir, self.PLATFORM, keyword)
        os.makedirs(result_dir, exist_ok=True)

        # 文件名使用时间戳
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"{timestamp}.json"
        filepath = os.path.join(result_dir, filename)

        # 保存结果
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(result, f, ensure_ascii=False, indent=2)

        return filepath


def main():
    """示例使用"""
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='小红书搜索推荐词接口')
    parser.add_argument(
        '--results-dir',
        type=str,
        default='data/search_recommendations',
        help='结果输出目录 (默认: data/search_recommendations)'
    )
    parser.add_argument(
        '--keyword',
        type=str,
        required=True,
        help='搜索关键词 (必填)'
    )
    args = parser.parse_args()

    # 创建API客户端实例
    client = XiaohongshuSearchRecommendations(results_dir=args.results_dir)

    # 获取推荐词并保存
    try:
        result = client.get_recommendations(args.keyword)
        filepath = client.save_result(args.keyword, result)
        print(f"Output: {filepath}")
    except Exception as e:
        print(f"Error: {e}", file=__import__('sys').stderr)


if __name__ == "__main__":
    main()