mirror of
https://github.com/AstrBotDevs/AstrBot
synced 2026-07-15 17:30:13 +08:00
ae29a7eaf9
- Add ExaWebSearchTool (web_search_exa) with keyword/semantic search, category filters, domain restrictions, and date range support - Add ExaGetContentsTool (exa_get_contents) for extracting web page content - Add _exa_search() and _exa_get_contents() API helpers hitting https://api.exa.ai/search and https://api.exa.ai/contents - Add _EXA_KEY_ROTATOR for multi-key rotation - Register Exa tools in _apply_web_search_tools() dispatch - Add Exa to WEB_SEARCH_CITATION_TOOL_NAMES for citation support - Add websearch_exa_key config default and provider option - Add i18n metadata for en-US, zh-CN, ru-RU - Add Exa section to docs (en + zh) - Add 6 unit tests covering search, contents, error handling, and legacy config migration Closes #5621
516 lines
15 KiB
Python
516 lines
15 KiB
Python
import json
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from astrbot.core.tools import web_search_tools as tools
|
|
|
|
|
|
class _FakeConfig(dict):
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.saved = False
|
|
|
|
def save_config(self):
|
|
self.saved = True
|
|
|
|
|
|
def test_normalize_legacy_web_search_config_migrates_firecrawl_key():
|
|
config = _FakeConfig(
|
|
{"provider_settings": {"websearch_firecrawl_key": "firecrawl-key"}}
|
|
)
|
|
|
|
tools.normalize_legacy_web_search_config(config)
|
|
|
|
assert config["provider_settings"]["websearch_firecrawl_key"] == ["firecrawl-key"]
|
|
assert config.saved is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_firecrawl_search_maps_web_results(monkeypatch):
|
|
async def fake_firecrawl_search(provider_settings, payload):
|
|
assert provider_settings["websearch_firecrawl_key"] == ["firecrawl-key"]
|
|
assert payload == {
|
|
"query": "AstrBot",
|
|
"limit": 3,
|
|
"sources": ["web"],
|
|
"country": "US",
|
|
}
|
|
return [
|
|
tools.SearchResult(
|
|
title="AstrBot",
|
|
url="https://example.com",
|
|
snippet="Search result",
|
|
)
|
|
]
|
|
|
|
monkeypatch.setattr(tools, "_firecrawl_search", fake_firecrawl_search)
|
|
tool = tools.FirecrawlWebSearchTool()
|
|
context = _context_with_provider_settings(
|
|
{"websearch_firecrawl_key": ["firecrawl-key"]}
|
|
)
|
|
|
|
result = await tool.call(context, query="AstrBot", limit=3, country="US")
|
|
|
|
assert json.loads(result)["results"] == [
|
|
{
|
|
"title": "AstrBot",
|
|
"url": "https://example.com",
|
|
"snippet": "Search result",
|
|
"index": json.loads(result)["results"][0]["index"],
|
|
}
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_firecrawl_search_maps_v2_data_list(monkeypatch):
|
|
session = _FakeFirecrawlSession(
|
|
_FakeFirecrawlResponse(
|
|
status=200,
|
|
json_data={
|
|
"success": True,
|
|
"data": [
|
|
{
|
|
"title": "AstrBot",
|
|
"url": "https://example.com",
|
|
"description": "Search result",
|
|
}
|
|
],
|
|
},
|
|
)
|
|
)
|
|
|
|
def fake_client_session(*, trust_env):
|
|
session.trust_env = trust_env
|
|
return session
|
|
|
|
monkeypatch.setattr(tools.aiohttp, "ClientSession", fake_client_session)
|
|
|
|
results = await tools._firecrawl_search(
|
|
{"websearch_firecrawl_key": ["firecrawl-key"]},
|
|
{"query": "AstrBot", "limit": 5, "sources": ["web"]},
|
|
)
|
|
|
|
assert session.posted == {
|
|
"url": "https://api.firecrawl.dev/v2/search",
|
|
"json": {"query": "AstrBot", "limit": 5, "sources": ["web"]},
|
|
"headers": {
|
|
"Authorization": "Bearer firecrawl-key",
|
|
"Content-Type": "application/json",
|
|
},
|
|
}
|
|
assert results == [
|
|
tools.SearchResult(
|
|
title="AstrBot", url="https://example.com", snippet="Search result"
|
|
)
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_firecrawl_search_maps_v2_grouped_web_data(monkeypatch):
|
|
session = _FakeFirecrawlSession(
|
|
_FakeFirecrawlResponse(
|
|
status=200,
|
|
json_data={
|
|
"success": True,
|
|
"data": {
|
|
"web": [
|
|
{
|
|
"title": "AstrBot",
|
|
"url": "https://example.com",
|
|
"description": "Search result",
|
|
}
|
|
]
|
|
},
|
|
},
|
|
)
|
|
)
|
|
|
|
def fake_client_session(*, trust_env):
|
|
session.trust_env = trust_env
|
|
return session
|
|
|
|
monkeypatch.setattr(tools.aiohttp, "ClientSession", fake_client_session)
|
|
|
|
results = await tools._firecrawl_search(
|
|
{"websearch_firecrawl_key": ["firecrawl-key"]},
|
|
{"query": "AstrBot", "limit": 5, "sources": ["web"]},
|
|
)
|
|
|
|
assert results == [
|
|
tools.SearchResult(
|
|
title="AstrBot", url="https://example.com", snippet="Search result"
|
|
)
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_firecrawl_search_payload_omits_tbs_and_uses_default_limit(monkeypatch):
|
|
async def fake_firecrawl_search(provider_settings, payload):
|
|
assert payload == {
|
|
"query": "AstrBot",
|
|
"limit": 5,
|
|
"sources": ["web"],
|
|
"country": "US",
|
|
}
|
|
return [
|
|
tools.SearchResult(
|
|
title="AstrBot",
|
|
url="https://example.com",
|
|
snippet="Search result",
|
|
)
|
|
]
|
|
|
|
monkeypatch.setattr(tools, "_firecrawl_search", fake_firecrawl_search)
|
|
tool = tools.FirecrawlWebSearchTool()
|
|
context = _context_with_provider_settings(
|
|
{"websearch_firecrawl_key": ["firecrawl-key"]}
|
|
)
|
|
|
|
result = await tool.call(
|
|
context,
|
|
query="AstrBot",
|
|
tbs="qdr:d",
|
|
country="US",
|
|
)
|
|
|
|
assert json.loads(result)["results"][0]["url"] == "https://example.com"
|
|
assert "tbs" not in tool.parameters["properties"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_firecrawl_extract_returns_scraped_markdown(monkeypatch):
|
|
async def fake_firecrawl_scrape(provider_settings, payload):
|
|
assert provider_settings["websearch_firecrawl_key"] == ["firecrawl-key"]
|
|
assert payload == {
|
|
"url": "https://example.com",
|
|
"formats": ["markdown"],
|
|
"onlyMainContent": True,
|
|
}
|
|
return {"url": "https://example.com", "markdown": "# Example"}
|
|
|
|
monkeypatch.setattr(tools, "_firecrawl_scrape", fake_firecrawl_scrape)
|
|
tool = tools.FirecrawlExtractWebPageTool()
|
|
context = _context_with_provider_settings(
|
|
{"websearch_firecrawl_key": ["firecrawl-key"]}
|
|
)
|
|
|
|
result = await tool.call(context, url="https://example.com")
|
|
|
|
assert result == "URL: https://example.com\nContent: # Example"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_firecrawl_search_uses_session_context(monkeypatch):
|
|
session = _FakeFirecrawlSession(
|
|
_FakeFirecrawlResponse(
|
|
status=200,
|
|
json_data={
|
|
"success": True,
|
|
"data": [
|
|
{
|
|
"title": "AstrBot",
|
|
"url": "https://example.com",
|
|
"description": "Search result",
|
|
}
|
|
],
|
|
},
|
|
)
|
|
)
|
|
|
|
def fake_client_session(*, trust_env):
|
|
session.trust_env = trust_env
|
|
return session
|
|
|
|
monkeypatch.setattr(tools.aiohttp, "ClientSession", fake_client_session)
|
|
|
|
await tools._firecrawl_search(
|
|
{"websearch_firecrawl_key": ["firecrawl-key"]},
|
|
{"query": "AstrBot"},
|
|
)
|
|
|
|
assert session.trust_env is True
|
|
assert session.entered is True
|
|
assert session.exited is True
|
|
assert session.posted == {
|
|
"url": "https://api.firecrawl.dev/v2/search",
|
|
"json": {"query": "AstrBot"},
|
|
"headers": {
|
|
"Authorization": "Bearer firecrawl-key",
|
|
"Content-Type": "application/json",
|
|
},
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_firecrawl_search_raises_error_for_http_errors(monkeypatch):
|
|
session = _FakeFirecrawlSession(
|
|
_FakeFirecrawlResponse(status=401, text_data="Unauthorized")
|
|
)
|
|
|
|
def fake_client_session(*, trust_env):
|
|
session.trust_env = trust_env
|
|
return session
|
|
|
|
monkeypatch.setattr(tools.aiohttp, "ClientSession", fake_client_session)
|
|
|
|
with pytest.raises(
|
|
Exception,
|
|
match="Firecrawl web search failed: Unauthorized, status: 401",
|
|
):
|
|
await tools._firecrawl_search(
|
|
{"websearch_firecrawl_key": ["firecrawl-key"]},
|
|
{"query": "AstrBot"},
|
|
)
|
|
|
|
assert session.trust_env is True
|
|
assert session.entered is True
|
|
assert session.exited is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_firecrawl_scrape_uses_request_setup(monkeypatch):
|
|
session = _FakeFirecrawlSession(
|
|
_FakeFirecrawlResponse(
|
|
status=200,
|
|
json_data={
|
|
"success": True,
|
|
"data": {"url": "https://example.com", "markdown": "# Example"},
|
|
},
|
|
)
|
|
)
|
|
|
|
def fake_client_session(*, trust_env):
|
|
session.trust_env = trust_env
|
|
return session
|
|
|
|
monkeypatch.setattr(tools.aiohttp, "ClientSession", fake_client_session)
|
|
|
|
result = await tools._firecrawl_scrape(
|
|
{"websearch_firecrawl_key": ["firecrawl-key"]},
|
|
{"url": "https://example.com", "formats": ["markdown"]},
|
|
)
|
|
|
|
assert result == {"url": "https://example.com", "markdown": "# Example"}
|
|
assert session.trust_env is True
|
|
assert session.entered is True
|
|
assert session.exited is True
|
|
assert session.posted == {
|
|
"url": "https://api.firecrawl.dev/v2/scrape",
|
|
"json": {"url": "https://example.com", "formats": ["markdown"]},
|
|
"headers": {
|
|
"Authorization": "Bearer firecrawl-key",
|
|
"Content-Type": "application/json",
|
|
},
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_firecrawl_scrape_raises_error_for_http_errors(monkeypatch):
|
|
session = _FakeFirecrawlSession(
|
|
_FakeFirecrawlResponse(status=401, text_data="Unauthorized")
|
|
)
|
|
|
|
def fake_client_session(*, trust_env):
|
|
session.trust_env = trust_env
|
|
return session
|
|
|
|
monkeypatch.setattr(tools.aiohttp, "ClientSession", fake_client_session)
|
|
|
|
with pytest.raises(
|
|
Exception,
|
|
match="Firecrawl web scraper failed: Unauthorized, status: 401",
|
|
):
|
|
await tools._firecrawl_scrape(
|
|
{"websearch_firecrawl_key": ["firecrawl-key"]},
|
|
{"url": "https://example.com", "formats": ["markdown"]},
|
|
)
|
|
|
|
assert session.trust_env is True
|
|
assert session.entered is True
|
|
assert session.exited is True
|
|
|
|
|
|
class _FakeFirecrawlResponse:
|
|
def __init__(self, status=200, json_data=None, text_data=""):
|
|
self.status = status
|
|
self.json_data = json_data or {}
|
|
self.text_data = text_data
|
|
|
|
async def __aenter__(self):
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc, tb):
|
|
return None
|
|
|
|
async def json(self):
|
|
return self.json_data
|
|
|
|
async def text(self):
|
|
return self.text_data
|
|
|
|
|
|
class _FakeFirecrawlSession:
|
|
def __init__(self, response):
|
|
self.response = response
|
|
self.trust_env = None
|
|
self.entered = False
|
|
self.exited = False
|
|
self.posted = None
|
|
|
|
async def __aenter__(self):
|
|
self.entered = True
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc, tb):
|
|
self.exited = True
|
|
return None
|
|
|
|
def post(self, url, json, headers):
|
|
self.posted = {"url": url, "json": json, "headers": headers}
|
|
return self.response
|
|
|
|
|
|
def _context_with_provider_settings(provider_settings):
|
|
config = {"provider_settings": provider_settings}
|
|
agent_context = SimpleNamespace(
|
|
context=SimpleNamespace(get_config=lambda umo: config),
|
|
event=SimpleNamespace(unified_msg_origin="test:private:session"),
|
|
)
|
|
return SimpleNamespace(context=agent_context)
|
|
|
|
|
|
# --- Exa tests ---
|
|
|
|
|
|
def test_normalize_legacy_web_search_config_migrates_exa_key():
|
|
config = _FakeConfig({"provider_settings": {"websearch_exa_key": "exa-key"}})
|
|
|
|
tools.normalize_legacy_web_search_config(config)
|
|
|
|
assert config["provider_settings"]["websearch_exa_key"] == ["exa-key"]
|
|
assert config.saved is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exa_search_maps_results(monkeypatch):
|
|
async def fake_exa_search(provider_settings, payload):
|
|
assert provider_settings["websearch_exa_key"] == ["exa-key"]
|
|
assert payload["query"] == "AstrBot"
|
|
assert payload["numResults"] == 5
|
|
return [
|
|
tools.SearchResult(
|
|
title="AstrBot",
|
|
url="https://example.com",
|
|
snippet="AI Agent Assistant",
|
|
)
|
|
]
|
|
|
|
monkeypatch.setattr(tools, "_exa_search", fake_exa_search)
|
|
tool = tools.ExaWebSearchTool()
|
|
context = _context_with_provider_settings({"websearch_exa_key": ["exa-key"]})
|
|
|
|
result = await tool.call(context, query="AstrBot", num_results=5)
|
|
|
|
parsed = json.loads(result)
|
|
assert parsed["results"][0]["title"] == "AstrBot"
|
|
assert parsed["results"][0]["url"] == "https://example.com"
|
|
assert parsed["results"][0]["snippet"] == "AI Agent Assistant"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exa_search_raw_api_call(monkeypatch):
|
|
session = _FakeFirecrawlSession(
|
|
_FakeFirecrawlResponse(
|
|
status=200,
|
|
json_data={
|
|
"results": [
|
|
{
|
|
"title": "AstrBot",
|
|
"url": "https://example.com",
|
|
"text": "AI Agent Assistant",
|
|
}
|
|
],
|
|
},
|
|
)
|
|
)
|
|
|
|
def fake_client_session(*, trust_env):
|
|
session.trust_env = trust_env
|
|
return session
|
|
|
|
monkeypatch.setattr(tools.aiohttp, "ClientSession", fake_client_session)
|
|
|
|
results = await tools._exa_search(
|
|
{"websearch_exa_key": ["exa-key"]},
|
|
{"query": "AstrBot", "numResults": 10, "type": "auto"},
|
|
)
|
|
|
|
assert session.posted["url"] == "https://api.exa.ai/search"
|
|
assert session.posted["headers"]["x-api-key"] == "exa-key"
|
|
assert results == [
|
|
tools.SearchResult(
|
|
title="AstrBot", url="https://example.com", snippet="AI Agent Assistant"
|
|
)
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exa_search_raises_on_http_error(monkeypatch):
|
|
session = _FakeFirecrawlSession(
|
|
_FakeFirecrawlResponse(status=401, text_data="Unauthorized")
|
|
)
|
|
|
|
def fake_client_session(*, trust_env):
|
|
session.trust_env = trust_env
|
|
return session
|
|
|
|
monkeypatch.setattr(tools.aiohttp, "ClientSession", fake_client_session)
|
|
|
|
with pytest.raises(
|
|
Exception,
|
|
match="Exa web search failed: Unauthorized, status: 401",
|
|
):
|
|
await tools._exa_search(
|
|
{"websearch_exa_key": ["exa-key"]},
|
|
{"query": "AstrBot"},
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exa_get_contents_returns_text(monkeypatch):
|
|
async def fake_exa_get_contents(provider_settings, payload):
|
|
assert provider_settings["websearch_exa_key"] == ["exa-key"]
|
|
assert payload["ids"] == ["https://example.com"]
|
|
return [{"url": "https://example.com", "text": "# Example Content"}]
|
|
|
|
monkeypatch.setattr(tools, "_exa_get_contents", fake_exa_get_contents)
|
|
tool = tools.ExaGetContentsTool()
|
|
context = _context_with_provider_settings({"websearch_exa_key": ["exa-key"]})
|
|
|
|
result = await tool.call(context, url="https://example.com")
|
|
|
|
assert result == "URL: https://example.com\nContent: # Example Content"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_exa_get_contents_raises_on_http_error(monkeypatch):
|
|
session = _FakeFirecrawlSession(
|
|
_FakeFirecrawlResponse(status=403, text_data="Forbidden")
|
|
)
|
|
|
|
def fake_client_session(*, trust_env):
|
|
session.trust_env = trust_env
|
|
return session
|
|
|
|
monkeypatch.setattr(tools.aiohttp, "ClientSession", fake_client_session)
|
|
|
|
with pytest.raises(
|
|
Exception,
|
|
match="Exa get contents failed: Forbidden, status: 403",
|
|
):
|
|
await tools._exa_get_contents(
|
|
{"websearch_exa_key": ["exa-key"]},
|
|
{"ids": ["https://example.com"]},
|
|
)
|