#!/usr/bin/env python3 """ 通用搜索工具 支持Google、Baidu、Bing等搜索引擎 """ import requests import json import os import argparse from datetime import datetime from typing import Dict, Any import sys from pathlib import Path # 添加项目根目录到路径并导入配置 sys.path.insert(0, str(Path(__file__).parent.parent.parent)) from lib.config import get_data_dir class CustomSearch: """通用搜索API封装类""" BASE_URL = "http://47.84.182.56:8001" TOOL_NAME = "custom_search" def __init__(self, results_dir: str = None): """ 初始化API客户端 Args: results_dir: 结果输出目录,默认为项目根目录下的 data/search 文件夹 """ self.api_url = f"{self.BASE_URL}/tools/call/{self.TOOL_NAME}" # 设置结果输出目录 if results_dir: self.results_base_dir = results_dir else: # 默认从配置读取 self.results_base_dir = get_data_dir("search") def search(self, keyword: str, platform: str = "google", timeout: int = 30) -> Dict[str, Any]: """ 执行搜索 Args: keyword: 搜索关键词 platform: 搜索平台,可选值:google, baidu, bing,默认为google timeout: 请求超时时间(秒),默认30秒 Returns: API响应的JSON数据 Raises: requests.exceptions.RequestException: 请求失败时抛出异常 """ payload = { "keyword": keyword, "platform": platform } try: response = requests.post( self.api_url, json=payload, timeout=timeout, headers={"Content-Type": "application/json"} ) response.raise_for_status() return response.json() except requests.exceptions.RequestException as e: print(f"请求失败: {e}") raise def save_result(self, keyword: str, platform: str, result: Dict[str, Any]) -> str: """ 保存结果到文件 目录结构: results/custom_search/平台/关键词/时间戳.json Args: keyword: 搜索关键词 platform: 搜索平台 result: API返回的结果 Returns: 保存的文件路径 """ # 创建目录结构: results/custom_search/平台/关键词/ result_dir = os.path.join(self.results_base_dir, "custom_search", 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=None, help='结果输出目录 (默认: 从配置读取)' ) parser.add_argument( '--keyword', type=str, required=True, help='搜索关键词 (必填)' ) parser.add_argument( '--platform', type=str, default='google', choices=['google', 'baidu', 'bing'], help='搜索平台 (默认: google)' ) args = parser.parse_args() # 创建API客户端实例 client = CustomSearch(results_dir=args.results_dir) # 执行搜索并保存 try: result = client.search(args.keyword, args.platform) filepath = client.save_result(args.keyword, args.platform, result) print(f"Output: {filepath}") except Exception as e: print(f"Error: {e}", file=__import__('sys').stderr) if __name__ == "__main__": main()