97 lines
2.9 KiB
Python
97 lines
2.9 KiB
Python
"""博查 (Bocha AI) Web Search API 客户端,共享给 web_search tool。"""
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
import httpx
|
|
import yaml
|
|
|
|
from core.paths import ROOT
|
|
|
|
_BOCHA_YAML = ROOT / "config" / "web" / "bocha.yaml"
|
|
|
|
|
|
class BochaError(RuntimeError):
|
|
"""博查 API 调用失败的统一异常。"""
|
|
|
|
|
|
@dataclass
|
|
class BochaConfig:
|
|
api_key: str
|
|
base_url: str
|
|
|
|
@classmethod
|
|
def load(cls, path: Optional[Path] = None) -> Optional["BochaConfig"]:
|
|
"""读 bocha.yaml + 解析 env 拿 api_key。
|
|
|
|
api_key env 未设 → 返 None(caller 据此决定是否注册 tool)。
|
|
yaml 不存在 → 返 None。
|
|
"""
|
|
p = path or _BOCHA_YAML
|
|
if not p.exists():
|
|
return None
|
|
data = yaml.safe_load(p.read_text(encoding="utf-8")) or {}
|
|
env = data.get("bocha_api_key_env") or "BOCHA_API_KEY"
|
|
key = os.environ.get(env, "").strip()
|
|
if not key:
|
|
return None
|
|
return cls(
|
|
api_key=key,
|
|
base_url=str(data.get("bocha_base_url") or "https://api.bochaai.com/v1").rstrip("/"),
|
|
)
|
|
|
|
|
|
class BochaClient:
|
|
"""轻量 httpx 封装: POST /v1/web-search, Bearer auth + 异常翻译。"""
|
|
|
|
def __init__(self, cfg: BochaConfig, timeout_s: float = 15.0) -> None:
|
|
self.cfg = cfg
|
|
self._client = httpx.Client(
|
|
base_url=cfg.base_url,
|
|
headers={
|
|
"Authorization": f"Bearer {cfg.api_key}",
|
|
"Content-Type": "application/json",
|
|
},
|
|
timeout=timeout_s,
|
|
)
|
|
|
|
def search(self, query: str, count: int = 10, freshness: str = "noLimit") -> dict:
|
|
"""调用博查 Web Search API,返回原始 dict。
|
|
|
|
freshness 可选: noLimit, oneDay, oneWeek, oneMonth, oneYear
|
|
count: 1-50
|
|
"""
|
|
body = {"query": query, "count": min(max(count, 1), 50), "freshness": freshness}
|
|
try:
|
|
resp = self._client.post("/web-search", json=body)
|
|
except httpx.TimeoutException as e:
|
|
raise BochaError(f"博查搜索超时: {e}") from e
|
|
except httpx.HTTPError as e:
|
|
raise BochaError(f"博查网络错误: {e}") from e
|
|
return self._parse(resp)
|
|
|
|
@staticmethod
|
|
def _parse(resp: httpx.Response) -> dict:
|
|
if resp.status_code >= 400:
|
|
try:
|
|
msg = resp.json().get("message", resp.text[:300])
|
|
except ValueError:
|
|
msg = resp.text[:300]
|
|
raise BochaError(f"博查 API → HTTP {resp.status_code}: {msg}")
|
|
try:
|
|
return resp.json()
|
|
except ValueError as e:
|
|
raise BochaError(f"博查 API → invalid JSON: {e}") from e
|
|
|
|
def close(self) -> None:
|
|
self._client.close()
|
|
|
|
def __enter__(self) -> "BochaClient":
|
|
return self
|
|
|
|
def __exit__(self, *_: object) -> None:
|
|
self.close()
|