273 lines
9.7 KiB
Python
273 lines
9.7 KiB
Python
"""通用搜索客户端（Serper / Exa / Tavily 路由）。
|
||
|
||
为 build_glossary.py 这类术语核查场景服务。
|
||
|
||
关键设计：
|
||
- `trust_env=False` 绕开系统 socks 代理（Clash on macOS 配 socks5 时 httpx 会 TLS EOF）
|
||
- 专利 / Scholar / News 优先 Serper，保证 Google Patents / Google Scholar 路径被真正调用
|
||
- 通用网页 Exa 优先，Tavily fallback
|
||
- 遇到配额问题自动降级或返回 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, *, strict_specialized: bool = True) -> None:
|
||
self._exa: ExaClient | None = None
|
||
self._tavily: TavilyClient | None = None
|
||
self._serper = None # 惰性实例化
|
||
self.strict_specialized = strict_specialized
|
||
try:
|
||
self._exa = ExaClient()
|
||
except SearchError:
|
||
pass
|
||
try:
|
||
self._tavily = TavilyClient()
|
||
except SearchError:
|
||
pass
|
||
self._has_serper_key = bool(os.environ.get("SERPER_API_KEY") or os.environ.get("SERPAPI_KEY"))
|
||
if not (self._exa or self._tavily or self._has_serper_key):
|
||
raise SearchError("no search API key available: set SERPER_API_KEY, SERPAPI_KEY, EXA_API_KEY, or TAVILY_API_KEY")
|
||
|
||
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 as exc:
|
||
if self.strict_specialized:
|
||
raise SearchError(f"serper patents failed: {exc}") from exc
|
||
# 降级：通用搜索加 site 限定
|
||
if self.strict_specialized:
|
||
raise SearchError("serper unavailable for patents route; refusing silent fallback")
|
||
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 as exc:
|
||
if self.strict_specialized:
|
||
raise SearchError(f"serper scholar failed: {exc}") from exc
|
||
if self.strict_specialized:
|
||
raise SearchError("serper unavailable for scholar route; refusing silent fallback")
|
||
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 as exc:
|
||
if self.strict_specialized:
|
||
raise SearchError(f"serper news failed: {exc}") from exc
|
||
if self.strict_specialized:
|
||
raise SearchError("serper unavailable for news route; refusing silent fallback")
|
||
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}")
|