mirror of
https://github.com/SMNETSTUDIO/WeChat-AI.git
synced 2026-08-20 18:13:43 +08:00
125 lines
3.9 KiB
Python
125 lines
3.9 KiB
Python
"""Web search backends (DDG default; optional SearXNG / Tavily)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from config import Settings
|
|
|
|
log = logging.getLogger("wechat-ai-tools.search")
|
|
|
|
|
|
class SearchError(RuntimeError):
|
|
pass
|
|
|
|
|
|
def _clamp_max(n: int | None, settings: Settings) -> int:
|
|
if n is None:
|
|
return max(1, min(settings.search_max_results, 10))
|
|
return max(1, min(int(n), settings.search_max_results, 10))
|
|
|
|
|
|
async def search_web(
|
|
query: str,
|
|
*,
|
|
max_results: int | None,
|
|
settings: Settings,
|
|
) -> list[dict[str, str]]:
|
|
q = (query or "").strip()
|
|
if not q:
|
|
raise SearchError("query is required")
|
|
if len(q) > 500:
|
|
raise SearchError("query too long (max 500)")
|
|
|
|
limit = _clamp_max(max_results, settings)
|
|
provider = (settings.search_provider or "ddg").strip().lower()
|
|
|
|
if provider == "tavily":
|
|
return await _search_tavily(q, limit, settings)
|
|
if provider == "searxng":
|
|
return await _search_searxng(q, limit, settings)
|
|
return await _search_ddg(q, limit, settings)
|
|
|
|
|
|
async def _search_ddg(query: str, limit: int, settings: Settings) -> list[dict[str, str]]:
|
|
try:
|
|
from duckduckgo_search import DDGS
|
|
except ImportError as exc:
|
|
raise SearchError("duckduckgo-search not installed") from exc
|
|
|
|
results: list[dict[str, str]] = []
|
|
timeout_s = max(1.0, settings.search_timeout_ms / 1000.0)
|
|
try:
|
|
# DDGS is sync; run in thread via anyio is ideal — use sync for MVP simplicity
|
|
with DDGS() as ddgs:
|
|
raw = list(ddgs.text(query, max_results=limit))
|
|
except Exception as exc: # noqa: BLE001
|
|
log.warning("ddg search failed: %s", type(exc).__name__)
|
|
raise SearchError(f"search failed: {type(exc).__name__}") from exc
|
|
|
|
for item in raw[:limit]:
|
|
results.append(
|
|
{
|
|
"title": str(item.get("title") or ""),
|
|
"url": str(item.get("href") or item.get("link") or ""),
|
|
"snippet": str(item.get("body") or item.get("snippet") or ""),
|
|
}
|
|
)
|
|
return results
|
|
|
|
|
|
async def _search_searxng(
|
|
query: str, limit: int, settings: Settings
|
|
) -> list[dict[str, str]]:
|
|
base = (settings.searxng_url or "").strip().rstrip("/")
|
|
if not base:
|
|
raise SearchError("SEARXNG_URL not configured")
|
|
timeout = max(1.0, settings.search_timeout_ms / 1000.0)
|
|
params = {"q": query, "format": "json", "categories": "general"}
|
|
async with httpx.AsyncClient(timeout=timeout) as client:
|
|
resp = await client.get(base, params=params)
|
|
resp.raise_for_status()
|
|
data: dict[str, Any] = resp.json()
|
|
out: list[dict[str, str]] = []
|
|
for item in (data.get("results") or [])[:limit]:
|
|
out.append(
|
|
{
|
|
"title": str(item.get("title") or ""),
|
|
"url": str(item.get("url") or ""),
|
|
"snippet": str(item.get("content") or item.get("snippet") or ""),
|
|
}
|
|
)
|
|
return out
|
|
|
|
|
|
async def _search_tavily(
|
|
query: str, limit: int, settings: Settings
|
|
) -> list[dict[str, str]]:
|
|
key = (settings.tavily_api_key or "").strip()
|
|
if not key:
|
|
raise SearchError("TAVILY_API_KEY not configured")
|
|
timeout = max(1.0, settings.search_timeout_ms / 1000.0)
|
|
payload = {
|
|
"api_key": key,
|
|
"query": query,
|
|
"max_results": limit,
|
|
"include_answer": False,
|
|
}
|
|
async with httpx.AsyncClient(timeout=timeout) as client:
|
|
resp = await client.post("https://api.tavily.com/search", json=payload)
|
|
resp.raise_for_status()
|
|
data: dict[str, Any] = resp.json()
|
|
out: list[dict[str, str]] = []
|
|
for item in (data.get("results") or [])[:limit]:
|
|
out.append(
|
|
{
|
|
"title": str(item.get("title") or ""),
|
|
"url": str(item.get("url") or ""),
|
|
"snippet": str(item.get("content") or ""),
|
|
}
|
|
)
|
|
return out
|