base_agent/tools/web_search.py

"""
tools/web_search.py
网络搜索工具 —— 支持 mock / SerpAPI / Brave Search
配置通过 settings.tools['web_search'] 读取
"""

import json
import time
from dataclasses import dataclass, field

from config.settings import settings
from tools.base_tool import BaseTool
from utils.logger import get_logger

logger = get_logger("TOOL.WebSearch")


def _cfg(key: str, fallback=None):
    return settings.tools['web_search'].get(key, fallback)


@dataclass
class SearchResult:
    title:   str
    url:     str
    snippet: str
    rank:    int = 0

    def __str__(self) -> str:
        return f"[{self.rank}] {self.title}\n    {self.url}\n    {self.snippet}"


class Tool(BaseTool):
    name        = "web_search"
    description = (
        "在互联网上搜索信息，返回相关网页的标题、链接和摘要。"
        "适用于需要实时信息、最新资讯或不确定的知识查询。"
    )
    parameters  = {
        "type": "object",
        "properties": {
            "query": {
                "type":        "string",
                "description": "搜索关键词或问题，例如: 'Python 3.12 新特性'",
            },
            "max_results": {
                "type":        "integer",
                "description": "返回结果数量（默认来自 config.yaml web_search.max_results）",
            },
        },
        "required": ["query"],
    }

    def execute(self, query: str = "", max_results: int | None = None, **_) -> str:
        if not query or not query.strip():
            return "❌ 参数错误: query 不能为空"

        n      = max_results or _cfg('max_results', 5)
        engine = _cfg('engine', 'mock')
        logger.info(
            f"🔍 搜索: {query}\n"
            f"   引擎={engine}  max_results={n}  "
            f"[config engine={_cfg('engine')} max_results={_cfg('max_results')}]"
        )

        match engine:
            case "serpapi":
                results = self._search_serpapi(query, n)
            case "brave":
                results = self._search_brave(query, n)
            case _:
                results = self._search_mock(query, n)

        if not results:
            return f"🔍 搜索 '{query}' 未找到相关结果"

        lines = [f"🔍 搜索结果: {query}  (共 {len(results)} 条)", "─" * 50]
        for r in results:
            lines.append(str(r))
        return "\n".join(lines)

    # ── 搜索引擎实现 ──────────────────────────────────────────

    @staticmethod
    def _search_mock(query: str, n: int) -> list[SearchResult]:
        """Mock 搜索（无需 API Key，用于测试）"""
        return [
            SearchResult(
                title=f"搜索结果 {i + 1}: {query}",
                url=f"https://example.com/result/{i + 1}",
                snippet=(
                    f"这是关于 '{query}' 的第 {i + 1} 条模拟搜索结果。"
                    f"实际使用请在 config.yaml 中配置 engine: serpapi 或 brave。"
                ),
                rank=i + 1,
            )
            for i in range(n)
        ]

    def _search_serpapi(self, query: str, n: int) -> list[SearchResult]:
        """SerpAPI 搜索"""
        api_key = _cfg('api_key', '')
        if not api_key:
            logger.warning("⚠️  SerpAPI api_key 未配置，回退到 mock 模式")
            return self._search_mock(query, n)
        try:
            import httpx
            timeout = _cfg('timeout', 10)
            resp = httpx.get(
                "https://serpapi.com/search",
                params={
                    "q":       query,
                    "num":     n,
                    "api_key": api_key,
                    "engine":  "google",
                },
                timeout=timeout,
            )
            resp.raise_for_status()
            data    = resp.json()
            organic = data.get("organic_results", [])
            return [
                SearchResult(
                    title=r.get("title", ""),
                    url=r.get("link", ""),
                    snippet=r.get("snippet", ""),
                    rank=i + 1,
                )
                for i, r in enumerate(organic[:n])
            ]
        except Exception as e:
            logger.error(f"❌ SerpAPI 搜索失败: {e}，回退到 mock")
            return self._search_mock(query, n)

    def _search_brave(self, query: str, n: int) -> list[SearchResult]:
        """Brave Search API"""
        api_key = _cfg('api_key', '')
        if not api_key:
            logger.warning("⚠️  Brave Search api_key 未配置，回退到 mock 模式")
            return self._search_mock(query, n)
        try:
            import httpx
            timeout = _cfg('timeout', 10)
            resp = httpx.get(
                "https://api.search.brave.com/res/v1/web/search",
                params={"q": query, "count": n},
                headers={
                    "Accept":               "application/json",
                    "Accept-Encoding":      "gzip",
                    "X-Subscription-Token": api_key,
                },
                timeout=timeout,
            )
            resp.raise_for_status()
            data    = resp.json()
            web     = data.get("web", {}).get("results", [])
            return [
                SearchResult(
                    title=r.get("title", ""),
                    url=r.get("url", ""),
                    snippet=r.get("description", ""),
                    rank=i + 1,
                )
                for i, r in enumerate(web[:n])
            ]
        except Exception as e:
            logger.error(f"❌ Brave Search 失败: {e}，回退到 mock")
            return self._search_mock(query, n)