Baidu Search

#!/usr/bin/env python3

# -*- coding: utf-8 -*-

"""百度 AI 搜索管理器"""

import requests

import json

import os

from datetime import datetime

from typing import List, Dict, Optional, Any





class BaiduSearchManager:

    """百度 AI 搜索管理器，支持网页搜索、百科、秒懂百科、AI 智能生成"""

    

    API_ENDPOINTS = {

        "web_search": "https://qianfan.baidubce.com/v2/ai_search/web_search",

        "baike": "https://appbuilder.baidu.com/v2/baike/lemma/get_content",

        "miaodong_baike": "https://appbuilder.baidu.com/v2/baike/second_know/search_video",

        "ai_chat": "https://qianfan.baidubce.com/v2/ai_search/chat/completions"

    }

    

    DEFAULT_API_TYPE = "web_search"

    

    def __init__(self, api_key: Optional[str] = None, config_file: Optional[str] = None):

        """初始化搜索管理器"""

        self.api_key = api_key

        

        # 未提供 api_key 时从配置文件加载

        if not self.api_key:

            # 默认配置文件路径

            if config_file is None:

                script_dir = os.path.dirname(os.path.abspath(__file__))

                config_file = os.path.join(script_dir, 'config.json')

            

            try:

                with open(config_file, 'r', encoding='utf-8') as f:

                    config = json.load(f)

                    self.api_key = config.get('api_key')

            except FileNotFoundError:

                raise ValueError(f"Config file not found: {config_file}. Please create it with your API key.")

            except json.JSONDecodeError:

                raise ValueError(f"Invalid JSON in config file: {config_file}")

        

        if not self.api_key:

            raise ValueError("API key is required. Set it in config.json or pass as parameter.")

        

        self.headers = {

            "X-Appbuilder-Authorization": f"Bearer {self.api_key}",

            "Content-Type": "application/json"

        }

    

    def get_current_date_context(self) -> Dict[str, Any]:

        """获取当前日期上下文"""

        now = datetime.now()

        return {

            'date': now.strftime('%Y-%m-%d'),

            'time': now.strftime('%H:%M:%S'),

            'weekday': now.strftime('%A'),

            'year': now.year,

            'month': now.month,

            'day': now.day,

        }

    

    def search(

        self,

        query: str,

        api_type: str = "web_search",

        max_results: int = 10,

        resource_types: Optional[List[Dict[str, Any]]] = None,

        sites: Optional[List[str]] = None,

        recency_filter: Optional[str] = None,

        safe_search: bool = False,

        include_date: bool = True

    ) -> Dict[str, Any]:

        """执行搜索"""

        # 校验 API 类型

        if api_type not in self.API_ENDPOINTS:

            raise ValueError(f"Invalid api_type: {api_type}. Must be one of {list(self.API_ENDPOINTS.keys())}")

        

        # 路由到对应搜索方法

        if api_type == "web_search":

            return self._search_web(query, max_results, resource_types, sites, recency_filter, safe_search, include_date)

        elif api_type == "baike":

            return self._search_baike(query, include_date)

        elif api_type == "miaodong_baike":

            return self._search_miaodong_baike(query, include_date)

        elif api_type == "ai_chat":

            return self._search_ai_chat(query, max_results, include_date)

    

    def _search_web(

        self,

        query: str,

        max_results: int,

        resource_types: Optional[List[Dict[str, Any]]],

        sites: Optional[List[str]],

        recency_filter: Optional[str],

        safe_search: bool,

        include_date: bool

    ) -> Dict[str, Any]:

        """网页搜索"""

        # 默认资源类型

        if resource_types is None:

            resource_types = [{"type": "web", "top_k": min(max_results, 50)}]

        

        # 构建请求

        payload = {

            "messages": [

                {

                    "role": "user",

                    "content": query[:72]  # 最长 72 字符