deep_research/scripts/lib/search_client.py

"""通用搜索客户端（Exa 优先，Tavily fallback）。

为 build_glossary.py 这类术语核查场景服务。

关键设计：
- `trust_env=False` 绕开系统 socks 代理（Clash on macOS 配 socks5 时 httpx 会 TLS EOF）
- Exa 优先：LinkedIn / 官网 / 百度百科返回质量最高
- 遇到配额问题自动降级到 Tavily 或返回 empty
- 不做深度 crawl，只要摘要
"""

from __future__ import annotations

import os
from dataclasses import dataclass
from typing import Any

import httpx


@dataclass
class SearchHit:
    title: str
    url: str
    snippet: str


class SearchError(RuntimeError):
    pass


class ExaClient:
    def __init__(self, api_key: str | None = None, timeout: float = 30.0) -> None:
        self.api_key = api_key or os.environ.get("EXA_API_KEY")
        if not self.api_key:
            raise SearchError("EXA_API_KEY not set")
        # trust_env=False 关键：不吃系统代理，避免 TLS EOF
        self._client = httpx.Client(trust_env=False, timeout=timeout)

    def close(self) -> None:
        self._client.close()

    def __enter__(self) -> "ExaClient":
        return self

    def __exit__(self, *_args: Any) -> None:
        self.close()

    def search(self, query: str, *, num_results: int = 5) -> list[SearchHit]:
        body = {
            "query": query,
            "numResults": num_results,
            "type": "auto",
            "contents": {"text": {"maxCharacters": 800}},
        }
        r = self._client.post(
            "https://api.exa.ai/search",
            json=body,
            headers={"x-api-key": self.api_key, "Content-Type": "application/json"},
        )
        if r.status_code != 200:
            raise SearchError(f"Exa HTTP {r.status_code}: {r.text[:200]}")
        data = r.json()
        out: list[SearchHit] = []
        for item in data.get("results", [])[:num_results]:
            out.append(
                SearchHit(
                    title=(item.get("title") or "")[:200],
                    url=item.get("url") or "",
                    snippet=(item.get("text") or item.get("snippet") or "")[:600],
                )
            )
        return out


class TavilyClient:
    def __init__(self, api_key: str | None = None, timeout: float = 30.0) -> None:
        self.api_key = api_key or os.environ.get("TAVILY_API_KEY")
        if not self.api_key:
            raise SearchError("TAVILY_API_KEY not set")
        self._client = httpx.Client(trust_env=False, timeout=timeout)

    def close(self) -> None:
        self._client.close()

    def __enter__(self) -> "TavilyClient":
        return self

    def __exit__(self, *_args: Any) -> None:
        self.close()

    def search(self, query: str, *, num_results: int = 5) -> list[SearchHit]:
        body = {
            "api_key": self.api_key,
            "query": query,
            "search_depth": "basic",
            "max_results": num_results,
            "include_answer": False,
            "include_raw_content": False,
        }
        r = self._client.post("https://api.tavily.com/search", json=body)
        if r.status_code != 200:
            raise SearchError(f"Tavily HTTP {r.status_code}: {r.text[:200]}")
        data = r.json()
        out: list[SearchHit] = []
        for item in data.get("results", [])[:num_results]:
            out.append(
                SearchHit(
                    title=(item.get("title") or "")[:200],
                    url=item.get("url") or "",
                    snippet=(item.get("content") or "")[:600],
                )
            )
        return out


class SearchClient:
    """统一搜索门面，支持多路由：

    - `search(query)`：通用网页搜索，优先 Exa → 降级 Tavily
    - `patents(query)`：专利检索，走 Serper（Google Patents）；失败则通用搜索补刀
    - `scholar(query)`：学术论文，走 Serper Scholar；失败则通用搜索补刀
    - `news(query)`：新闻检索，走 Serper News；失败则通用搜索补刀

    所有客户端都延迟导入 serper_client，避免没装 SERPAPI_KEY 时 import 炸。
    """

    def __init__(self) -> None:
        self._exa: ExaClient | None = None
        self._tavily: TavilyClient | None = None
        self._serper = None  # 惰性实例化
        try:
            self._exa = ExaClient()
        except SearchError:
            pass
        try:
            self._tavily = TavilyClient()
        except SearchError:
            pass
        if not (self._exa or self._tavily):
            raise SearchError(
                "neither EXA_API_KEY nor TAVILY_API_KEY available"
            )

    def _get_serper(self):
        """惰性创建 SerperClient。没 key 时返回 None。"""
        if self._serper is False:
            return None
        if self._serper is None:
            try:
                from scripts.lib.serper_client import SerperClient
                self._serper = SerperClient()
            except Exception:
                self._serper = False
                return None
        return self._serper

    def close(self) -> None:
        if self._exa:
            self._exa.close()
        if self._tavily:
            self._tavily.close()
        if self._serper and self._serper is not False:
            self._serper.close()

    def __enter__(self) -> "SearchClient":
        return self

    def __exit__(self, *_args: Any) -> None:
        self.close()

    def search(self, query: str, *, num_results: int = 5) -> list[SearchHit]:
        """通用网页搜索。Exa 首选，Tavily 备选。"""
        if self._exa:
            try:
                return self._exa.search(query, num_results=num_results)
            except SearchError:
                pass
        if self._tavily:
            try:
                return self._tavily.search(query, num_results=num_results)
            except SearchError:
                pass
        return []

    def patents(self, query: str, *, num_results: int = 10) -> list[SearchHit]:
        """专利检索：Serper 走 Google Patents 最准。降级到通用搜索 + site 限定。"""
        serper = self._get_serper()
        if serper:
            try:
                hits = serper.patents(query, num_results=num_results)
                return [SearchHit(h.title, h.url, h.snippet) for h in hits]
            except Exception:
                pass
        # 降级：通用搜索加 site 限定
        return self.search(f"site:patents.google.com {query}", num_results=num_results)

    def scholar(
        self,
        query: str,
        *,
        num_results: int = 10,
        year_low: int | None = None,
    ) -> list[SearchHit]:
        """学术论文：Serper Scholar 带引用数。降级到通用搜索。"""
        serper = self._get_serper()
        if serper:
            try:
                hits = serper.scholar(query, num_results=num_results, year_low=year_low)
                return [
                    SearchHit(
                        title=h.title,
                        url=h.url,
                        snippet=f"{h.snippet}  | {h.source} | 引用 {h.cited_by}" if h.cited_by else h.snippet,
                    )
                    for h in hits
                ]
            except Exception:
                pass
        return self.search(query, num_results=num_results)

    def news(
        self,
        query: str,
        *,
        num_results: int = 10,
        time_range: str | None = None,
    ) -> list[SearchHit]:
        """新闻检索：Serper News。降级到通用搜索。"""
        serper = self._get_serper()
        if serper:
            try:
                hits = serper.news(query, num_results=num_results, time_range=time_range)
                return [
                    SearchHit(
                        title=h.title,
                        url=h.url,
                        snippet=f"{h.snippet}  | {h.source} | {h.date}" if h.date else h.snippet,
                    )
                    for h in hits
                ]
            except Exception:
                pass
        return self.search(query, num_results=num_results)


if __name__ == "__main__":
    from scripts.lib.zenmux_client import load_secrets
    load_secrets()
    with SearchClient() as c:
        print("--- 通用: Mabwell 迈威生物 ---")
        for h in c.search("Mabwell 迈威生物 biopharmaceutical", num_results=3):
            print(f"  {h.title[:80]}")
            print(f"    {h.url}")
        print("\n--- 专利: dual-target siRNA ---")
        for h in c.patents("dual-target siRNA GalNAc", num_results=3):
            print(f"  {h.title[:80]}")
            print(f"    {h.url}")
        print("\n--- Scholar: dual-target RNAi 2024 ---")
        for h in c.scholar("dual-target RNAi drug", num_results=3, year_low=2023):
            print(f"  {h.title[:80]}")
            print(f"    {h.url}")