322 lines
12 KiB
Python
322 lines
12 KiB
Python
"""BaiduSearchTool - 百度搜索工具，支持优雅降级
|
||
|
||
通过百度搜索 API 执行关键词搜索，返回搜索结果列表。
|
||
当百度搜索 API 不可用时，返回包含降级提示的错误信息。
|
||
"""
|
||
|
||
import json
|
||
import logging
|
||
import urllib.parse
|
||
|
||
import httpx
|
||
|
||
from agentkit.tools.base import Tool
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class BaiduSearchTool(Tool):
|
||
"""百度搜索工具 - 执行关键词搜索，返回搜索结果
|
||
|
||
支持两种模式：
|
||
1. 百度搜索 API（需要 API key 配置）
|
||
2. 直接抓取百度搜索结果页（降级模式，无需 API key）
|
||
|
||
当两种模式都不可用时，返回包含降级提示的错误信息。
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
name: str = "baidu_search",
|
||
description: str = "执行百度搜索，返回搜索结果列表",
|
||
input_schema: dict[str, object] | None = None,
|
||
output_schema: dict[str, object] | None = None,
|
||
version: str = "1.0.0",
|
||
tags: list[str] | None = None,
|
||
api_key: str | None = None,
|
||
api_url: str | None = None,
|
||
):
|
||
super().__init__(
|
||
name=name,
|
||
description=description,
|
||
input_schema=input_schema or self._default_input_schema(),
|
||
output_schema=output_schema or self._default_output_schema(),
|
||
version=version,
|
||
tags=tags or ["search", "baidu"],
|
||
)
|
||
self._api_key = api_key
|
||
self._api_url = api_url
|
||
|
||
@staticmethod
|
||
def _default_input_schema() -> dict[str, object]:
|
||
return {
|
||
"type": "object",
|
||
"properties": {
|
||
"query": {
|
||
"type": "string",
|
||
"description": "搜索关键词",
|
||
},
|
||
"max_results": {
|
||
"type": "integer",
|
||
"description": "最大返回结果数",
|
||
"default": 5,
|
||
},
|
||
},
|
||
"required": ["query"],
|
||
}
|
||
|
||
@staticmethod
|
||
def _default_output_schema() -> dict[str, object]:
|
||
return {
|
||
"type": "object",
|
||
"properties": {
|
||
"results": {
|
||
"type": "array",
|
||
"items": {
|
||
"type": "object",
|
||
"properties": {
|
||
"title": {"type": "string"},
|
||
"url": {"type": "string"},
|
||
"snippet": {"type": "string"},
|
||
},
|
||
},
|
||
"description": "搜索结果列表",
|
||
},
|
||
"total": {"type": "integer", "description": "结果总数"},
|
||
"success": {"type": "boolean", "description": "是否成功"},
|
||
"error": {"type": "string", "description": "错误信息（仅失败时）"},
|
||
},
|
||
}
|
||
|
||
async def execute(self, **kwargs) -> dict:
|
||
"""执行百度搜索
|
||
|
||
Args:
|
||
query: 搜索关键词（必需）
|
||
max_results: 最大返回结果数（默认 5）
|
||
|
||
Returns:
|
||
包含 results 列表和 success 布尔值的字典
|
||
"""
|
||
query = kwargs.get("query")
|
||
if not query:
|
||
return {"error": "query 参数是必需的", "results": [], "total": 0, "success": False}
|
||
|
||
max_results = kwargs.get("max_results", 5)
|
||
|
||
# 优先使用 API 模式
|
||
if self._api_key and self._api_url:
|
||
return await self._search_via_api(query, max_results)
|
||
|
||
# 降级：直接抓取百度搜索结果页
|
||
return await self._search_via_scrape(query, max_results)
|
||
|
||
async def _search_via_api(self, query: str, max_results: int) -> dict:
|
||
"""通过百度搜索 API 执行搜索"""
|
||
try:
|
||
params = {
|
||
"query": query,
|
||
"num": max_results,
|
||
}
|
||
url = f"{self._api_url}?{urllib.parse.urlencode(params)}"
|
||
async with httpx.AsyncClient(timeout=30) as client:
|
||
resp = await client.get(
|
||
url,
|
||
headers={
|
||
"User-Agent": "AgentKit/1.0",
|
||
"Authorization": f"Bearer {self._api_key}",
|
||
},
|
||
)
|
||
resp.raise_for_status()
|
||
data = resp.json()
|
||
|
||
results = []
|
||
for item in data.get("results", [])[:max_results]:
|
||
results.append({
|
||
"title": item.get("title", ""),
|
||
"url": item.get("url", ""),
|
||
"snippet": item.get("snippet", ""),
|
||
})
|
||
|
||
return {"results": results, "total": len(results), "success": True}
|
||
|
||
except Exception as e:
|
||
logger.error(f"BaiduSearchTool API 搜索失败: {e}")
|
||
# 降级到抓取模式
|
||
return await self._search_via_scrape(query, max_results)
|
||
|
||
async def _search_via_scrape(self, query: str, max_results: int) -> dict:
|
||
"""通过直接抓取百度搜索结果页执行搜索（降级模式）"""
|
||
try:
|
||
encoded_query = urllib.parse.quote(query)
|
||
url = f"https://www.baidu.com/s?wd={encoded_query}&rn={max_results}"
|
||
async with httpx.AsyncClient(timeout=30, follow_redirects=True) as client:
|
||
resp = await client.get(
|
||
url,
|
||
headers={
|
||
"User-Agent": (
|
||
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "
|
||
"AppleWebKit/537.36 (KHTML, like Gecko) "
|
||
"Chrome/131.0.0.0 Safari/537.36"
|
||
),
|
||
"Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
|
||
"Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8",
|
||
"Accept-Encoding": "gzip, deflate, br",
|
||
"Connection": "keep-alive",
|
||
"Cache-Control": "max-age=0",
|
||
"Sec-Fetch-Dest": "document",
|
||
"Sec-Fetch-Mode": "navigate",
|
||
"Sec-Fetch-Site": "none",
|
||
"Sec-Fetch-User": "?1",
|
||
"Upgrade-Insecure-Requests": "1",
|
||
},
|
||
)
|
||
html = resp.text
|
||
|
||
# Check if we got a captcha page
|
||
if "验证" in html and len(html) < 5000:
|
||
logger.warning("Baidu returned captcha page, search unavailable")
|
||
return {
|
||
"error": "Baidu search blocked by captcha",
|
||
"results": [],
|
||
"total": 0,
|
||
"success": False,
|
||
}
|
||
|
||
# 简单解析搜索结果（基于百度搜索结果页 HTML 结构）
|
||
results = self._parse_baidu_html(html, max_results)
|
||
|
||
if not results:
|
||
# Try alternative parsing
|
||
results = self._parse_baidu_html_alt(html, max_results)
|
||
|
||
return {"results": results, "total": len(results), "success": True}
|
||
|
||
except Exception as e:
|
||
logger.error(f"BaiduSearchTool 抓取搜索失败: {e}")
|
||
return {
|
||
"error": f"百度搜索不可用: {e}",
|
||
"results": [],
|
||
"total": 0,
|
||
"success": False,
|
||
}
|
||
|
||
@staticmethod
|
||
def _parse_baidu_html(html: str, max_results: int) -> list[dict[str, str]]:
|
||
"""解析百度搜索结果页 HTML，提取标题、URL、摘要
|
||
|
||
注意：百度 HTML 结构可能变化，此解析器尽力提取关键信息。
|
||
"""
|
||
import re
|
||
|
||
results: list[dict[str, str]] = []
|
||
|
||
# 匹配百度搜索结果块 - multiple patterns for different Baidu page versions
|
||
# Pattern 1: <h3 class="t"> with href
|
||
pattern1 = re.compile(
|
||
r'<h3[^>]*class="[^"]*t[^"]*"[^>]*>.*?href="([^"]*)"[^>]*>(.*?)</a>',
|
||
re.DOTALL,
|
||
)
|
||
# Pattern 2: <h3> with data-url or inside <div class="result">
|
||
pattern2 = re.compile(
|
||
r'<h3[^>]*>.*?<a[^>]*href="([^"]*)"[^>]*>(.*?)</a>',
|
||
re.DOTALL,
|
||
)
|
||
# Snippet patterns
|
||
snippet_pattern1 = re.compile(
|
||
r'<span[^>]*class="[^"]*content-right_[^"]*"[^>]*>(.*?)</span>',
|
||
re.DOTALL,
|
||
)
|
||
snippet_pattern2 = re.compile(
|
||
r'<div[^>]*class="[^"]*c-abstract[^"]*"[^>]*>(.*?)</div>',
|
||
re.DOTALL,
|
||
)
|
||
snippet_pattern3 = re.compile(
|
||
r'<span[^>]*class="[^"]*content-right_[^"]*"[^>]*>(.*?)</span>',
|
||
re.DOTALL,
|
||
)
|
||
|
||
# Try pattern 1 first
|
||
for match in pattern1.finditer(html):
|
||
if len(results) >= max_results:
|
||
break
|
||
url = match.group(1)
|
||
title = re.sub(r"<[^>]+>", "", match.group(2)).strip()
|
||
if not title or len(title) < 2:
|
||
continue
|
||
# Skip Baidu internal links that aren't redirect links
|
||
if "baidu.com" in url and "baidu.com/link?" not in url:
|
||
continue
|
||
if not url.startswith("http") and "baidu.com/link?" not in url:
|
||
continue
|
||
|
||
snippet = ""
|
||
for sp in [snippet_pattern1, snippet_pattern2, snippet_pattern3]:
|
||
snippet_match = sp.search(html[match.end():match.end() + 2000])
|
||
if snippet_match:
|
||
snippet = re.sub(r"<[^>]+>", "", snippet_match.group(1)).strip()
|
||
if snippet:
|
||
break
|
||
|
||
results.append({
|
||
"title": title[:200],
|
||
"url": url,
|
||
"snippet": snippet[:300] if snippet else "",
|
||
})
|
||
|
||
# If pattern 1 found nothing, try pattern 2
|
||
if not results:
|
||
for match in pattern2.finditer(html):
|
||
if len(results) >= max_results:
|
||
break
|
||
url = match.group(1)
|
||
title = re.sub(r"<[^>]+>", "", match.group(2)).strip()
|
||
if not title or len(title) < 2:
|
||
continue
|
||
if "baidu.com" in url and "baidu.com/link?" not in url:
|
||
continue
|
||
if not url.startswith("http") and "baidu.com/link?" not in url:
|
||
continue
|
||
|
||
snippet = ""
|
||
for sp in [snippet_pattern1, snippet_pattern2, snippet_pattern3]:
|
||
snippet_match = sp.search(html[match.end():match.end() + 2000])
|
||
if snippet_match:
|
||
snippet = re.sub(r"<[^>]+>", "", snippet_match.group(1)).strip()
|
||
if snippet:
|
||
break
|
||
|
||
results.append({
|
||
"title": title[:200],
|
||
"url": url,
|
||
"snippet": snippet[:300] if snippet else "",
|
||
})
|
||
|
||
return results
|
||
|
||
@staticmethod
|
||
def _parse_baidu_html_alt(html: str, max_results: int) -> list[dict[str, str]]:
|
||
"""Alternative Baidu HTML parser - broader pattern matching."""
|
||
import re
|
||
|
||
results: list[dict[str, str]] = []
|
||
|
||
# Generic pattern: any <a> tag with baidu.com/link redirect
|
||
pattern = re.compile(
|
||
r'<a[^>]*href="(https?://www\.baidu\.com/link\?[^"]*)"[^>]*>(.*?)</a>',
|
||
re.DOTALL,
|
||
)
|
||
for match in pattern.finditer(html):
|
||
if len(results) >= max_results:
|
||
break
|
||
url = match.group(1)
|
||
title = re.sub(r"<[^>]+>", "", match.group(2)).strip()
|
||
if title and len(title) > 2:
|
||
results.append({
|
||
"title": title[:200],
|
||
"url": url,
|
||
"snippet": "",
|
||
})
|
||
|
||
return results
|