mirror of
https://github.com/HKUDS/nanobot.git
synced 2026-05-19 16:12:30 +00:00
fix(web): back off Brave search rate limits
This commit is contained in:
parent
0f3677c0d8
commit
b2ac609bb5
@ -272,23 +272,37 @@ class WebSearchTool(Tool):
|
||||
logger.warning("BRAVE_API_KEY not set, falling back to DuckDuckGo")
|
||||
return await self._search_duckduckgo(query, n)
|
||||
try:
|
||||
async with httpx.AsyncClient(proxy=self.proxy) as client:
|
||||
r = await client.get(
|
||||
"https://api.search.brave.com/res/v1/web/search",
|
||||
params={"q": query, "count": n},
|
||||
headers = {
|
||||
"Accept": "application/json",
|
||||
"X-Subscription-Token": api_key,
|
||||
"User-Agent": self.user_agent,
|
||||
},
|
||||
}
|
||||
async with httpx.AsyncClient(proxy=self.proxy) as client:
|
||||
for attempt in range(2):
|
||||
r = await client.get(
|
||||
"https://api.search.brave.com/res/v1/web/search",
|
||||
params={"q": query, "count": n},
|
||||
headers=headers,
|
||||
timeout=10.0,
|
||||
)
|
||||
if r.status_code != 429:
|
||||
break
|
||||
if attempt == 0:
|
||||
logger.warning("Brave search rate limited; retrying once in 1.0s")
|
||||
await asyncio.sleep(1.0)
|
||||
r.raise_for_status()
|
||||
items = [
|
||||
{"title": x.get("title", ""), "url": x.get("url", ""), "content": x.get("description", "")}
|
||||
for x in r.json().get("web", {}).get("results", [])
|
||||
]
|
||||
return _format_results(query, items, n)
|
||||
except httpx.HTTPStatusError as e:
|
||||
if e.response.status_code == 429:
|
||||
return (
|
||||
"Error: Brave search rate limited after retry. "
|
||||
"Retry later or reduce consecutive web_search calls."
|
||||
)
|
||||
return f"Error: {e}"
|
||||
except Exception as e:
|
||||
return f"Error: {e}"
|
||||
|
||||
|
||||
@ -19,7 +19,10 @@ def _tool(
|
||||
)
|
||||
|
||||
|
||||
def _response(status: int = 200, json: dict | None = None) -> httpx.Response:
|
||||
def _response(
|
||||
status: int = 200,
|
||||
json: dict | None = None,
|
||||
) -> httpx.Response:
|
||||
"""Build a mock httpx.Response with a dummy request attached."""
|
||||
r = httpx.Response(status, json=json)
|
||||
r._request = httpx.Request("GET", "https://mock")
|
||||
@ -62,6 +65,55 @@ async def test_brave_search(monkeypatch):
|
||||
assert "https://example.com" in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_brave_search_retries_rate_limit_once(monkeypatch):
|
||||
calls = {"n": 0}
|
||||
sleeps: list[float] = []
|
||||
|
||||
async def mock_sleep(delay: float):
|
||||
sleeps.append(delay)
|
||||
|
||||
async def mock_get(self, url, **kw):
|
||||
calls["n"] += 1
|
||||
if calls["n"] == 1:
|
||||
return _response(status=429, json={"error": "rate limit"})
|
||||
return _response(json={
|
||||
"web": {"results": [{"title": "Recovered", "url": "https://example.com", "description": "ok"}]}
|
||||
})
|
||||
|
||||
monkeypatch.setattr("nanobot.agent.tools.web.asyncio.sleep", mock_sleep)
|
||||
monkeypatch.setattr(httpx.AsyncClient, "get", mock_get)
|
||||
|
||||
tool = _tool(provider="brave", api_key="brave-key")
|
||||
result = await tool.execute(query="nanobot", count=1)
|
||||
|
||||
assert calls["n"] == 2
|
||||
assert "Recovered" in result
|
||||
assert sleeps == [1.0]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_brave_search_returns_clear_rate_limit_after_retries(monkeypatch):
|
||||
calls = {"n": 0}
|
||||
|
||||
async def mock_sleep(delay: float):
|
||||
return None
|
||||
|
||||
async def mock_get(self, url, **kw):
|
||||
calls["n"] += 1
|
||||
return _response(status=429, json={"error": "rate limit"})
|
||||
|
||||
monkeypatch.setattr("nanobot.agent.tools.web.asyncio.sleep", mock_sleep)
|
||||
monkeypatch.setattr(httpx.AsyncClient, "get", mock_get)
|
||||
|
||||
tool = _tool(provider="brave", api_key="brave-key")
|
||||
result = await tool.execute(query="nanobot", count=1)
|
||||
|
||||
assert calls["n"] == 2
|
||||
assert "Brave search rate limited" in result
|
||||
assert "consecutive web_search" in result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tavily_search(monkeypatch):
|
||||
async def mock_post(self, url, **kw):
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user