实现了完整的网页搜索与抓取内置工具集: 1. 新增搜索工具链:支持SearXNG与Tavily多Provider降级、滑动窗口限流、统一结果格式 2. 新增网页抓取工具:单页抓取、站点递归抓取、内容入库知识库功能 3. 完善安全防护:SSRF校验、robots.txt合规、请求限流 4. 自动注册工具到系统注册表,无需额外配置即可使用
165 lines
5.9 KiB
Python
165 lines
5.9 KiB
Python
"""搜索 Provider 协议与 SearXNG/Tavily 实现。"""
|
||
import os
|
||
from typing import Protocol
|
||
|
||
import httpx
|
||
from tavily import AsyncTavilyClient
|
||
|
||
from yuxi.config import config
|
||
|
||
from .models import SearchResult
|
||
|
||
|
||
class SearchProvider(Protocol):
|
||
"""搜索 Provider 协议,所有实现必须符合此契约。"""
|
||
|
||
@property
|
||
def name(self) -> str:
|
||
"""Provider 名称,如 'searxng' / 'tavily'。"""
|
||
...
|
||
|
||
def is_available(self) -> bool:
|
||
"""Provider 是否可用(每次调用时检查,非启动时固定)。
|
||
|
||
- SearXNG:检查 config.web_search_searxng_base_url 非空
|
||
- Tavily:检查 config.web_search_tavily_api_key 非空或 TAVILY_API_KEY 环境变量存在
|
||
"""
|
||
...
|
||
|
||
async def search(self, query: str, *, max_results: int = 5, **opts) -> list[SearchResult]:
|
||
"""执行搜索,返回结构化结果列表。
|
||
|
||
异常语义:
|
||
- 网络错误 / 429 / 5xx:抛异常,由 ProviderChain 决定是否降级
|
||
- 4xx(非 429):抛异常,ProviderChain 不降级(视为配置错误)
|
||
- 空结果:返回空列表,不抛异常,ProviderChain 不降级
|
||
"""
|
||
...
|
||
|
||
|
||
class SearXNGProvider:
|
||
"""SearXNG 元搜索引擎 Provider,HTTP 调用自部署实例。"""
|
||
|
||
name = "searxng"
|
||
|
||
def __init__(self):
|
||
self._client: httpx.AsyncClient | None = None # 懒初始化,复用连接池
|
||
|
||
def is_available(self) -> bool:
|
||
"""检查 config.web_search_searxng_base_url 是否非空。"""
|
||
return bool(config.web_search_searxng_base_url)
|
||
|
||
async def search(self, query: str, *, max_results: int = 5, **opts) -> list[SearchResult]:
|
||
"""调用 SearXNG /search 接口,返回结构化结果。"""
|
||
base_url = config.web_search_searxng_base_url
|
||
timeout = config.web_search_searxng_timeout
|
||
params = {
|
||
"q": query,
|
||
"format": "json",
|
||
"pageno": 1,
|
||
"safesearch": config.web_search_searxng_safesearch,
|
||
}
|
||
if config.web_search_searxng_engines:
|
||
params["engines"] = ",".join(config.web_search_searxng_engines)
|
||
language = config.web_search_searxng_language
|
||
if language and language != "auto":
|
||
params["language"] = language
|
||
|
||
client = self._get_client()
|
||
response = await client.get(
|
||
f"{base_url}/search", params=params, timeout=timeout
|
||
)
|
||
response.raise_for_status() # 4xx/5xx 抛异常,由 chain 判断是否降级
|
||
data = response.json()
|
||
results = []
|
||
for item in data.get("results", [])[:max_results]:
|
||
results.append(
|
||
SearchResult(
|
||
title=item.get("title", ""),
|
||
url=item.get("url", ""),
|
||
snippet=(item.get("content") or "")[:500],
|
||
published_at=item.get("publishedDate"),
|
||
source_provider=f"searxng:{item['engine']}" if item.get("engine") else "searxng",
|
||
)
|
||
)
|
||
return results
|
||
|
||
def _get_client(self) -> httpx.AsyncClient:
|
||
"""懒初始化 httpx.AsyncClient,复用连接池。"""
|
||
if self._client is None or self._client.is_closed:
|
||
self._client = httpx.AsyncClient()
|
||
return self._client
|
||
|
||
|
||
class TavilyProvider:
|
||
"""Tavily 搜索 Provider,直接使用 tavily-python SDK。
|
||
|
||
不再经 langchain_tavily.TavilySearch 包装,原因:
|
||
1. TavilySearch 是 LangChain 工具实例,返回格式与 Provider 协议不一致
|
||
2. 包装层冗余且难以控制参数
|
||
3. tavily-python>=0.7.0 已在 pyproject.toml 中
|
||
"""
|
||
|
||
name = "tavily"
|
||
|
||
def __init__(self):
|
||
self._client: AsyncTavilyClient | None = None # 懒初始化,复用 SDK 客户端
|
||
|
||
def is_available(self) -> bool:
|
||
"""检查 config.web_search_tavily_api_key 或 TAVILY_API_KEY 环境变量。"""
|
||
api_key = config.web_search_tavily_api_key or os.getenv("TAVILY_API_KEY")
|
||
return bool(api_key)
|
||
|
||
async def search(self, query: str, *, max_results: int = 5, **opts) -> list[SearchResult]:
|
||
"""调用 Tavily SDK,返回结构化结果。"""
|
||
api_key = config.web_search_tavily_api_key or os.getenv("TAVILY_API_KEY")
|
||
if not api_key:
|
||
raise RuntimeError("Tavily API Key 未配置")
|
||
client = self._get_client(api_key)
|
||
response = await client.search(
|
||
query=query,
|
||
max_results=max_results,
|
||
timeout=config.web_search_tavily_timeout,
|
||
)
|
||
results = []
|
||
for item in response.get("results", []):
|
||
results.append(
|
||
SearchResult(
|
||
title=item.get("title", ""),
|
||
url=item.get("url", ""),
|
||
snippet=(item.get("content") or "")[:500],
|
||
published_at=item.get("published_date"),
|
||
source_provider="tavily",
|
||
)
|
||
)
|
||
return results
|
||
|
||
def _get_client(self, api_key: str) -> AsyncTavilyClient:
|
||
"""懒初始化 AsyncTavilyClient,复用 SDK 客户端。
|
||
|
||
api_key 变更时重建客户端。
|
||
"""
|
||
if self._client is None or getattr(self._client, "api_key", None) != api_key:
|
||
self._client = AsyncTavilyClient(api_key=api_key)
|
||
return self._client
|
||
|
||
|
||
_searxng_provider: SearXNGProvider | None = None
|
||
_tavily_provider: TavilyProvider | None = None
|
||
|
||
|
||
def get_searxng_provider() -> SearXNGProvider:
|
||
"""获取 SearXNGProvider 单例,复用 httpx 连接池。"""
|
||
global _searxng_provider
|
||
if _searxng_provider is None:
|
||
_searxng_provider = SearXNGProvider()
|
||
return _searxng_provider
|
||
|
||
|
||
def get_tavily_provider() -> TavilyProvider:
|
||
"""获取 TavilyProvider 单例,复用 SDK 客户端。"""
|
||
global _tavily_provider
|
||
if _tavily_provider is None:
|
||
_tavily_provider = TavilyProvider()
|
||
return _tavily_provider
|