crawler.py

"""
模块 3: 爬虫/搜索 API (Crawler & Search)

封装搜索 API 调用 + 网页内容抓取。
支持多种搜索后端 (SerpAPI / Bing / 自建) 和统一的搜索结果格式。
"""

from __future__ import annotations

import asyncio
import re
from dataclasses import dataclass, field
from typing import Optional
from urllib.parse import quote_plus

import httpx
from bs4 import BeautifulSoup

from config import ModuleConfig


# ── 统一数据结构 ────────────────────────────────────────────────

@dataclass
class SearchResult:
    """单条搜索结果"""
    title: str
    url: str
    snippet: str
    source: str = ""           # 来源域名
    content: str = ""          # 抓取到的详情页正文 (可选)
    fetch_error: str = ""      # 抓取失败原因
    
    @property
    def display(self) -> str:
        """格式化输出"""
        parts = [f"📄 {self.title}", f"   URL: {self.url}", f"   摘要: {self.snippet}"]
        if self.content:
            parts.append(f"   正文: {self.content[:300]}...")
        if self.fetch_error:
            parts.append(f"   ⚠️ 抓取失败: {self.fetch_error}")
        return "\n".join(parts)


@dataclass
class SearchResponse:
    """搜索响应"""
    query: str
    results: list[SearchResult] = field(default_factory=list)
    total_results: int = 0
    search_time: float = 0.0
    error: str = ""
    
    @property
    def context_text(self) -> str:
        """将搜索结果拼接为 LLM 可用的上下文字符串"""
        parts = []
        for i, r in enumerate(self.results, 1):
            part = f"[{i}] {r.title}\n来源: {r.url}\n摘要: {r.snippet}"
            if r.content:
                part += f"\n详细内容:\n{r.content[:1000]}"
            parts.append(part)
        return "\n\n---\n\n".join(parts)


# ── HTML 清洗 ───────────────────────────────────────────────────

def _clean_html(html: str, url: str = "") -> str:
    """从 HTML 中提取正文，去除导航、广告、脚本等噪音"""
    try:
        soup = BeautifulSoup(html, "html.parser")
    except Exception:
        return ""
    
    # 移除噪音标签
    for tag in soup.find_all(["script", "style", "nav", "footer", "header",
                               "aside", "noscript", "iframe", "form"]):
        tag.decompose()
    
    # 移除常见的广告/侧边栏 class
    noise_classes = [
        "sidebar", "advertisement", "ad-", "banner", "popup",
        "nav-", "menu", "footer", "header-", "comment",
        "social", "share", "related", "recommend",
    ]
    for tag in soup.find_all(class_=re.compile("|".join(noise_classes))):
        tag.decompose()
    
    # 优先从常见内容区域提取
    content_selectors = [
        "article", "main", '[role="main"]',
        ".post-content", ".article-content", ".entry-content",
        "#content", "#article", ".content",
        ".post-body", ".markdown-body",
    ]
    
    content_tag = None
    for selector in content_selectors:
        content_tag = soup.select_one(selector)
        if content_tag:
            break
    
    if content_tag is None:
        content_tag = soup.body or soup
    
    text = content_tag.get_text(separator="\n", strip=True)
    
    # 压缩多余空行
    text = re.sub(r"\n{3,}", "\n\n", text)
    text = re.sub(r"[ \t]{3,}", "  ", text)
    
    # 截断过长内容 (节省 token)
    if len(text) > 3000:
        text = text[:3000] + "\n...[内容截断]"
    
    return text


# ── 搜索 API 适配器 ─────────────────────────────────────────────

async def _search_serpapi(query: str, config: ModuleConfig) -> SearchResponse:
    """SerpAPI 搜索 (底层用 Google)"""
    
    params = {
        "q": query,
        "api_key": config.search.api_key,
        "engine": "google",
        "num": config.search.max_results,
        "hl": "zh-CN",
    }
    
    async with httpx.AsyncClient(timeout=config.search.fetch_timeout + 5) as client:
        resp = await client.get(config.search.base_url, params=params)
        resp.raise_for_status()
        data = resp.json()
    
    results = []
    for item in data.get("organic_results", [])[: config.search.max_results]:
        results.append(SearchResult(
            title=item.get("title", ""),
            url=item.get("link", ""),
            snippet=item.get("snippet", ""),
            source=item.get("source", ""),
        ))
    
    return SearchResponse(
        query=query,
        results=results,
        total_results=data.get("search_information", {}).get("total_results", 0),
    )


async def _search_bing(query: str, config: ModuleConfig) -> SearchResponse:
    """Bing Web Search API"""
    
    async with httpx.AsyncClient(timeout=config.search.fetch_timeout + 5) as client:
        resp = await client.get(
            f"{config.search.base_url}/v7.0/search",
            headers={"Ocp-Apim-Subscription-Key": config.search.api_key},
            params={
                "q": query,
                "count": config.search.max_results,
                "mkt": "zh-CN",
                "textFormat": "Raw",
            },
        )
        resp.raise_for_status()
        data = resp.json()
    
    results = []
    for item in data.get("webPages", {}).get("value", [])[: config.search.max_results]:
        results.append(SearchResult(
            title=item.get("name", ""),
            url=item.get("url", ""),
            snippet=item.get("snippet", ""),
        ))
    
    return SearchResponse(
        query=query,
        results=results,
        total_results=data.get("webPages", {}).get("totalEstimatedMatches", 0),
    )


async def _search_bing_free(query: str, config: ModuleConfig) -> SearchResponse:
    """Bing 直接抓取 — 免 API Key，零成本。冷门实体名自动回退。"""
    import re
    from bs4 import BeautifulSoup
    
    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
        "Accept-Language": "zh-CN,zh;q=0.9",
    }
    
    async def _do_search(q):
        url = f"{config.search.base_url}?q={q}&setlang=zh-Hans&count={config.search.max_results}"
        async with httpx.AsyncClient(timeout=15, follow_redirects=True) as client:
            resp = await client.get(url, headers=headers)
            resp.raise_for_status()
        soup = BeautifulSoup(resp.text, "html.parser")
        results = []
        for li in soup.select("li.b_algo"):
            h2 = li.find("h2")
            a = h2.find("a") if h2 else li.find("a")
            p = li.select_one(".b_caption p, .b_lineclamp2")
            if a:
                results.append(SearchResult(
                    title=a.get_text(strip=True),
                    url=a.get("href", ""),
                    snippet=p.get_text(strip=True) if p else "",
                ))
        return SearchResponse(query=q, results=results[: config.search.max_results])
    
    response = await _do_search(query)
    
    # 冷门实体名检测: 标题不含核心词 → 回退搜裸名或双词
    entities = re.findall(r'[\u4e00-\u9fff]{2,4}', query)
    if entities and len(response.results) > 3:
        main = entities[0]  # 取第一个中文词 (通常是主语实体名，而非修饰词)
        hit = sum(1 for r in response.results if main in r.title)
        if hit <= len(response.results) * 0.3:
            # 回退策略：先用双词(实体+修饰)，不行再用裸名
            fallback_queries = []
            if len(entities) > 1:
                fallback_queries.append(f"{main} {entities[1]}")
            fallback_queries.append(main)
            
            for fbq in fallback_queries:
                if fbq == query:
                    continue
                if config.verbose:
                    print(f"[Crawler] 实体 '{main}' 匹配率 {hit}/{len(response.results)}，回退搜: {fbq}")
                fallback = await _do_search(fbq)
                if fallback.results:
                    response = fallback
                    response.query = query
                    break
    
    return response


# ── 网页抓取 ────────────────────────────────────────────────────

async def _fetch_page_content(url: str, config: ModuleConfig) -> str:
    """抓取单个网页的正文内容"""
    try:
        headers = {
            "User-Agent": (
                "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                "AppleWebKit/537.36 (KHTML, like Gecko) "
                "Chrome/120.0.0.0 Safari/537.36"
            ),
            "Accept": "text/html,application/xhtml+xml",
            "Accept-Language": "zh-CN,zh;q=0.9,en;q=0.8",
        }
        async with httpx.AsyncClient(
            timeout=config.search.fetch_timeout,
            follow_redirects=True,
            headers=headers,
        ) as client:
            resp = await client.get(url)
            resp.raise_for_status()
            return _clean_html(resp.text, url)
    except Exception as e:
        return f""


async def _search_360(query: str, config: ModuleConfig) -> SearchResponse:
    """360 搜索直接抓取 — 免费，中文分词优秀"""
    from urllib.parse import quote
    
    url = f"https://www.so.com/s?q={quote(query)}&pn=1"
    headers = {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
        "Accept-Language": "zh-CN,zh;q=0.9",
    }
    
    async with httpx.AsyncClient(timeout=15, follow_redirects=True) as client:
        resp = await client.get(url, headers=headers)
        resp.raise_for_status()
    
    soup = BeautifulSoup(resp.text, "html.parser")
    results = []
    for item in soup.select(".result, .res-list"):
        a = item.select_one("h3 a")
        p = item.select_one(".res-desc, .res-rich, .res-summary")
        if a:
            results.append(SearchResult(
                title=a.get_text(strip=True),
                url=a.get("href", ""),
                snippet=p.get_text(strip=True) if p else "",
            ))
    
    return SearchResponse(
        query=query,
        results=results[: config.search.max_results],
    )


async def _enrich_results(
    response: SearchResponse,
    config: ModuleConfig,
) -> SearchResponse:
    """并发抓取搜索结果详情页内容 (只抓前 N 条)"""
    top_n = config.search.fetch_content
    if not top_n or not response.results:
        return response
    
    targets = response.results[:top_n]
    if config.verbose:
        print(f"[Crawler] 开始并发抓取前 {len(targets)} 个页面...")
    
    async def fetch_one(result: SearchResult) -> SearchResult:
        if not result.url:
            return result
        try:
            result.content = await _fetch_page_content(result.url, config)
        except Exception as e:
            result.fetch_error = str(e)[:200]
        return result
    
    tasks = [fetch_one(r) for r in targets]
    enriched = await asyncio.gather(*tasks)
    # 替换回原列表
    response.results[:top_n] = enriched
    
    if config.verbose:
        fetched = sum(1 for r in enriched if r.content)
        print(f"[Crawler] 成功抓取 {fetched}/{len(targets)} 个页面")
    
    return response


# ── 对外接口 ────────────────────────────────────────────────────

SEARCH_PROVIDERS = {
    "360": _search_360,
    "bing_free": _search_bing_free,
    "serpapi": _search_serpapi,
    "bing": _search_bing,
}


async def search_and_fetch(
    query: str,
    config: ModuleConfig,
) -> SearchResponse:
    """
    执行搜索并可选抓取详情页内容。
    
    Args:
        query: 搜索查询词 (应该是 rewrite 后的)
        config: 模块配置
    """
    provider = config.search.provider
    search_fn = SEARCH_PROVIDERS.get(provider)
    
    if search_fn is None:
        return SearchResponse(
            query=query,
            error=f"不支持的搜索后端: {provider}. 可选: {list(SEARCH_PROVIDERS.keys())}",
        )
    
    if config.verbose:
        print(f"[Crawler] 使用 {provider} 搜索: {query}")
    
    try:
        response = await search_fn(query, config)
    except Exception as e:
        return SearchResponse(query=query, error=f"搜索失败: {e}")
    
    if config.verbose:
        print(f"[Crawler] 搜索完成, {len(response.results)} 条结果 ({response.total_results} total)")
    
    # 可选: 抓取详情页
    if config.search.fetch_content and response.results:
        response = await _enrich_results(response, config)
    
    return response