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:
    """统一搜索门面：先用 Exa，失败/配额问题降级 Tavily。"""

    def __init__(self) -> None:
        self._exa: ExaClient | None = None
        self._tavily: TavilyClient | None = 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 close(self) -> None:
        if self._exa:
            self._exa.close()
        if self._tavily:
            self._tavily.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
        if self._exa:
            try:
                return self._exa.search(query, num_results=num_results)
            except SearchError as e:
                msg = str(e).lower()
                if "exceed" in msg or "quota" in msg or "429" in msg or "402" in msg:
                    # 降级
                    pass
                else:
                    # 其它错误继续往下试
                    pass
        if self._tavily:
            try:
                return self._tavily.search(query, num_results=num_results)
            except SearchError:
                pass
        return []


if __name__ == "__main__":
    from scripts.lib.zenmux_client import load_secrets
    load_secrets()
    with SearchClient() as c:
        hits = c.search("Mabwell 迈威生物 biopharmaceutical", num_results=3)
        for i, h in enumerate(hits, 1):
            print(f"[{i}] {h.title[:80]}")
            print(f"    {h.url}")
            print(f"    {h.snippet[:160]}")