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")
|
logger.warning("BRAVE_API_KEY not set, falling back to DuckDuckGo")
|
||||||
return await self._search_duckduckgo(query, n)
|
return await self._search_duckduckgo(query, n)
|
||||||
try:
|
try:
|
||||||
|
headers = {
|
||||||
|
"Accept": "application/json",
|
||||||
|
"X-Subscription-Token": api_key,
|
||||||
|
"User-Agent": self.user_agent,
|
||||||
|
}
|
||||||
async with httpx.AsyncClient(proxy=self.proxy) as client:
|
async with httpx.AsyncClient(proxy=self.proxy) as client:
|
||||||
r = await client.get(
|
for attempt in range(2):
|
||||||
"https://api.search.brave.com/res/v1/web/search",
|
r = await client.get(
|
||||||
params={"q": query, "count": n},
|
"https://api.search.brave.com/res/v1/web/search",
|
||||||
headers={
|
params={"q": query, "count": n},
|
||||||
"Accept": "application/json",
|
headers=headers,
|
||||||
"X-Subscription-Token": api_key,
|
timeout=10.0,
|
||||||
"User-Agent": self.user_agent,
|
)
|
||||||
},
|
if r.status_code != 429:
|
||||||
timeout=10.0,
|
break
|
||||||
)
|
if attempt == 0:
|
||||||
|
logger.warning("Brave search rate limited; retrying once in 1.0s")
|
||||||
|
await asyncio.sleep(1.0)
|
||||||
r.raise_for_status()
|
r.raise_for_status()
|
||||||
items = [
|
items = [
|
||||||
{"title": x.get("title", ""), "url": x.get("url", ""), "content": x.get("description", "")}
|
{"title": x.get("title", ""), "url": x.get("url", ""), "content": x.get("description", "")}
|
||||||
for x in r.json().get("web", {}).get("results", [])
|
for x in r.json().get("web", {}).get("results", [])
|
||||||
]
|
]
|
||||||
return _format_results(query, items, n)
|
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:
|
except Exception as e:
|
||||||
return f"Error: {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."""
|
"""Build a mock httpx.Response with a dummy request attached."""
|
||||||
r = httpx.Response(status, json=json)
|
r = httpx.Response(status, json=json)
|
||||||
r._request = httpx.Request("GET", "https://mock")
|
r._request = httpx.Request("GET", "https://mock")
|
||||||
@ -62,6 +65,55 @@ async def test_brave_search(monkeypatch):
|
|||||||
assert "https://example.com" in result
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_tavily_search(monkeypatch):
|
async def test_tavily_search(monkeypatch):
|
||||||
async def mock_post(self, url, **kw):
|
async def mock_post(self, url, **kw):
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user